From 5848c500e6a8e4e6c76f2c1ec085f1056ec03f03 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 24 Aug 2026 22:27:23 +0800 Subject: [PATCH] fix(persistence): avoid duplicate index dumps --- .../components/file_graph/local_file_graph.py | 9 +++++ reme/components/file_graph/nx_file_graph.py | 9 +++++ reme/components/keyword_index/bm25_index.py | 15 ++++++++ tests/unit/test_file_store_consistency.py | 36 +++++++++++++++++++ 4 files changed, 69 insertions(+) diff --git a/reme/components/file_graph/local_file_graph.py b/reme/components/file_graph/local_file_graph.py index f47c75e4..9ffbe363 100644 --- a/reme/components/file_graph/local_file_graph.py +++ b/reme/components/file_graph/local_file_graph.py @@ -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 ------------------------------------------------------- diff --git a/reme/components/file_graph/nx_file_graph.py b/reme/components/file_graph/nx_file_graph.py index 27257925..c2918194 100644 --- a/reme/components/file_graph/nx_file_graph.py +++ b/reme/components/file_graph/nx_file_graph.py @@ -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 ------------------------------------------------------- diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 0cb3db86..c3f7a713 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -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 diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 6c49b4ab..224776aa 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -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."""