From 65f3bc3ef28b43d2a18e656393660adb96f203db Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 26 Aug 2026 16:08:50 +0800 Subject: [PATCH] fix(persistence): align dumps with component ownership --- .../components/file_graph/local_file_graph.py | 9 --- reme/components/file_graph/nx_file_graph.py | 9 --- .../file_store/faiss_local_file_store.py | 6 +- .../components/file_store/local_file_store.py | 16 +++-- .../file_store/zvec_local_file_store.py | 6 +- reme/components/keyword_index/bm25_index.py | 15 ----- tests/unit/test_file_graph.py | 62 +++++++++++++++++++ tests/unit/test_file_store_consistency.py | 6 +- tests/unit/test_keyword_index.py | 22 +++++++ 9 files changed, 105 insertions(+), 46 deletions(-) diff --git a/reme/components/file_graph/local_file_graph.py b/reme/components/file_graph/local_file_graph.py index 9ffbe363..f47c75e4 100644 --- a/reme/components/file_graph/local_file_graph.py +++ b/reme/components/file_graph/local_file_graph.py @@ -19,7 +19,6 @@ 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 --------------------------------------------------------- @@ -36,17 +35,13 @@ 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}") @@ -83,8 +78,6 @@ 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) @@ -103,7 +96,6 @@ 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) @@ -127,7 +119,6 @@ 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 c2918194..27257925 100644 --- a/reme/components/file_graph/nx_file_graph.py +++ b/reme/components/file_graph/nx_file_graph.py @@ -25,7 +25,6 @@ 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 --------------------------------------------------------- @@ -35,20 +34,16 @@ 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}") @@ -74,8 +69,6 @@ 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): @@ -87,7 +80,6 @@ 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: @@ -109,7 +101,6 @@ 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/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 2e02bfe2..06fe6eda 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -496,10 +496,10 @@ class FaissLocalFileStore(LocalFileStore): self.faiss_idmap_path.unlink(missing_ok=True) return False - async def dump(self) -> None: - """Persist chunks JSONL via the parent, then write the FAISS sidecar atomically.""" + async def _dump_owned_state(self) -> None: + """Persist chunks and the FAISS sidecar, excluding dependency snapshots.""" async with self._faiss_dump_lock: - await super().dump() + await super()._dump_owned_state() if self._faiss_index is None or self.embedding_store is None: return try: diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 3bb797f8..37bdb0fc 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -100,7 +100,11 @@ class LocalFileStore(BaseFileStore): async def _close(self) -> None: self._closing = True await self._cancel_embedding_backfill() - await self.dump() + # Dependencies are closed separately by Application (reverse + # topological order) or BaseComponent (owned standalone dependencies). + # Persist only this store's local state here so each component writes + # exactly once during shutdown. + await self._dump_owned_state() self.file_chunks.clear() await super()._close() @@ -529,9 +533,8 @@ class LocalFileStore(BaseFileStore): elapsed = time.monotonic() - started_at self.logger.info(f"{self.name}: keyword index rebuild complete: total={total}, elapsed={elapsed:.2f}s") - async def dump(self) -> None: - """Atomically rewrite the JSONL, then cascade dump into keyword_index and file_graph.""" - assert self.file_graph is not None + async def _dump_owned_state(self) -> None: + """Persist state owned by this store, excluding dependency snapshots.""" try: write_jsonl_zst( self.chunks_path, @@ -541,6 +544,11 @@ class LocalFileStore(BaseFileStore): self.logger.info(f"Saved {len(self.file_chunks)} chunks to {self.chunks_path}") except Exception as e: self.logger.exception(f"Failed to write {self.chunks_path}: {e}") + + async def dump(self) -> None: + """Persist a complete store/index/graph consistency checkpoint.""" + assert self.file_graph is not None + await self._dump_owned_state() if self.keyword_index: await self.keyword_index.dump() await self.file_graph.dump() diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 4b234a3e..5925fd84 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -319,9 +319,9 @@ class ZvecLocalFileStore(LocalFileStore): f"elapsed={time.monotonic() - started_at:.3f}s", ) - async def dump(self) -> None: - """Persist chunks JSONL via the parent, then flush zvec and write the sidecar.""" - await super().dump() + async def _dump_owned_state(self) -> None: + """Persist chunks and zvec state, excluding dependency snapshots.""" + await super()._dump_owned_state() if self._collection is None or self.embedding_store is None: return try: diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index c3f7a713..0cb3db86 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -52,10 +52,6 @@ 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 ----------------------------------------------------------- @@ -249,8 +245,6 @@ 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] = [] @@ -278,8 +272,6 @@ 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 = {} @@ -362,11 +354,8 @@ 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) @@ -374,7 +363,6 @@ 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}") @@ -388,7 +376,6 @@ 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}") @@ -407,7 +394,6 @@ class BM25Index(BaseKeywordIndex): self._posting_tfs = {} self._idf_cache = {} self.index_file.unlink(missing_ok=True) - self._needs_persist = False # -- Compaction ----------------------------------------------------------- @@ -493,4 +479,3 @@ 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_graph.py b/tests/unit/test_file_graph.py index a125d3fe..9c53c92f 100644 --- a/tests/unit/test_file_graph.py +++ b/tests/unit/test_file_graph.py @@ -329,6 +329,68 @@ def test_rebuild_links_idempotent(backend_cls): asyncio.run(run()) +def test_nx_rebuild_links_repair_survives_restart(): + """A repaired NetworkX snapshot must not revert to stale serialized edges.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + graph = NxFileGraph() + await graph.start() + await graph.upsert_nodes( + [ + make_node("a.md", [("c.md", None)]), + make_node("b.md"), + make_node("c.md"), + ], + ) + await graph.dump() + + # Model an old/torn snapshot: the node payload says a -> b while + # the separately serialized NetworkX edge still says a -> c. + graph._graph.nodes["a.md"]["node"] = make_node("a.md", [("b.md", None)]) + await graph.dump() + await graph.close() + + repaired = NxFileGraph() + await repaired.start() + assert {link.target_path for link in await repaired.get_outlinks("a.md")} == {"c.md"} + await repaired.rebuild_links() + assert {link.target_path for link in await repaired.get_outlinks("a.md")} == {"b.md"} + await repaired.close() + + reopened = NxFileGraph() + await reopened.start() + assert {link.target_path for link in await reopened.get_outlinks("a.md")} == {"b.md"} + await reopened.close() + + asyncio.run(run()) + + +def test_local_load_merged_state_survives_restart(): + """Loading a snapshot into a modified local graph must not suppress its later dump.""" + + async def run(): + with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir): + seed = LocalFileGraph() + await seed.start() + await seed.upsert_nodes([make_node("a.md")]) + await seed.close() + + merged = LocalFileGraph() + await merged.start() + await merged.upsert_nodes([make_node("b.md")]) + await merged.load() + assert {node.path for node in await merged.get_nodes()} == {"a.md", "b.md"} + await merged.close() + + reopened = LocalFileGraph() + await reopened.start() + assert {node.path for node in await reopened.get_nodes()} == {"a.md", "b.md"} + await reopened.close() + + asyncio.run(run()) + + @pytest.mark.parametrize("backend_cls", BACKENDS) def test_persistence_roundtrip(backend_cls): """close() dumps; a fresh instance loads the same nodes from disk.""" diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 224776aa..d7fdda97 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -285,8 +285,8 @@ 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().""" +def test_close_persists_each_component_once(monkeypatch): + """Store and owned dependency shutdown each persist their own state once.""" bm25_writes = 0 graph_writes = 0 real_pickle_dump = bm25_index_module.pickle.dump @@ -1433,7 +1433,7 @@ def test_faiss_rejects_stale_sidecar_after_partial_dump(): # t2: simulate a crash between the two writes in dump(): only the # parent's JSONL write lands; the sidecar stays at the alpha # generation. (No close() -- the process is presumed dead.) - await LocalFileStore.dump(store_a) + await LocalFileStore._dump_owned_state(store_a) # t3: restart. The stale sidecar must be rejected by the digest. store_b = _new_faiss_store("t_faiss_stale_sidecar") diff --git a/tests/unit/test_keyword_index.py b/tests/unit/test_keyword_index.py index 607a9e36..b9554e5d 100644 --- a/tests/unit/test_keyword_index.py +++ b/tests/unit/test_keyword_index.py @@ -666,6 +666,28 @@ def test_dump_load_roundtrip_preserves_state(): run(go()) +def test_runtime_parameter_update_survives_restart(): + """Persist k1/b updates made after loading an existing index.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + seed = await create_bm25() + await seed.add_docs({"d1": "hello world"}) + await seed.close() + + changed = await create_bm25() + assert (changed.k1, changed.b) == (1.5, 0.75) + changed.k1 = 2.0 + changed.b = 0.4 + await changed.close() + + reopened = await create_bm25() + assert (reopened.k1, reopened.b) == (2.0, 0.4) + await reopened.close() + + run(go()) + + def test_load_missing_file_keeps_empty_state(): """Calling load() with no file on disk is a no-op."""