diff --git a/reme/application.py b/reme/application.py index b4b0cb69..ab06c2be 100644 --- a/reme/application.py +++ b/reme/application.py @@ -15,6 +15,7 @@ from .plugin import resolve_plugin_runtime from .schema import ComponentConfig, Response, StreamChunk from .utils import execute_stream_task, print_logo, get_logger +from .utils.async_utils import complete_in_thread T = TypeVar("T", bound=BaseComponent) _NodeKey = tuple[str, str] @@ -229,8 +230,9 @@ async def _close_started_components(self) -> None: self.logger.exception(f"Failed to close {component_type_name(c.component_type)}:{c.name}: {e}") self._started_components.clear() if self.context.thread_pool is not None: - self.context.thread_pool.shutdown(wait=True) + thread_pool = self.context.thread_pool self.context.thread_pool = None + await complete_in_thread(thread_pool.shutdown, True) async def update_component(self, component_enum: ComponentType, name: str, /, **kwargs) -> BaseComponent: """Update an existing component by type/name; never creates missing components.""" diff --git a/reme/components/file_catalog/local_file_catalog.py b/reme/components/file_catalog/local_file_catalog.py index 1cfe1199..2681a87c 100644 --- a/reme/components/file_catalog/local_file_catalog.py +++ b/reme/components/file_catalog/local_file_catalog.py @@ -5,6 +5,7 @@ from .base_file_catalog import BaseFileCatalog from ..component_registry import R from ...schema import FileNode +from ...utils.async_utils import complete_in_thread from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -24,12 +25,18 @@ async def load(self) -> None: async with self._io_lock: if not self._catalog_file.exists(): return - await self._read_jsonl() + loaded = await complete_in_thread( + self._read_jsonl_sync, + self._catalog_file, + self.encoding, + self._nodes, + ) + self._nodes = loaded self.logger.debug(f"Loaded {len(self._nodes)} nodes from {self._catalog_file}") async def dump(self) -> None: async with self._io_lock: - await self._write_jsonl() + await complete_in_thread(self._write_jsonl_sync) self.logger.info(f"Saved {len(self._nodes)} nodes to {self._catalog_file}") async def upsert(self, nodes: list[FileNode]) -> None: @@ -49,11 +56,16 @@ async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]: return list(self._nodes.values()) return [self._nodes[p] for p in paths if p in self._nodes] - async def _read_jsonl(self) -> None: - for line in read_jsonl_zst(self._catalog_file, self.encoding): + @staticmethod + def _read_jsonl_sync(path, encoding: str, existing: dict[str, FileNode]) -> dict[str, FileNode]: + """Read, decompress, and parse a catalog checkpoint off-loop.""" + nodes = dict(existing) + for line in read_jsonl_zst(path, encoding): if stripped := line.strip(): node = FileNode.model_validate_json(stripped) - self._nodes[node.path] = node + nodes[node.path] = node + return nodes - async def _write_jsonl(self) -> None: + def _write_jsonl_sync(self) -> None: + """Serialize, compress, and atomically publish the locked catalog state.""" write_jsonl_zst(self._catalog_file, (n.model_dump_json() for n in self._nodes.values()), self.encoding) diff --git a/reme/components/file_graph/local_file_graph.py b/reme/components/file_graph/local_file_graph.py index f47c75e4..b871d56d 100644 --- a/reme/components/file_graph/local_file_graph.py +++ b/reme/components/file_graph/local_file_graph.py @@ -1,11 +1,13 @@ """Pure-Python file-graph backend (no external deps).""" +import asyncio from pathlib import Path from .base_file_graph import BaseFileGraph from ..component_registry import R from ...enumeration import LinkScopeEnum from ...schema import FileLink, FileNode +from ...utils.async_utils import complete_in_thread from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst @@ -18,6 +20,7 @@ def __init__(self, **kwargs): self._nodes: dict[str, FileNode] = {} self._inverse: dict[str, set[str]] = {} # real target → sources self._pending: dict[str, set[str]] = {} # virtual target → sources + self._io_lock = asyncio.Lock() self._graph_file: Path = self.component_metadata_path / f"{self.name}.jsonl.zst" # -- Lifecycle --------------------------------------------------------- @@ -25,26 +28,56 @@ def __init__(self, **kwargs): async def _start(self) -> None: self.component_metadata_path.mkdir(parents=True, exist_ok=True) await super()._start() # base calls load() - await self.rebuild_links() async def load(self) -> None: - if not self._graph_file.exists(): - return - try: - for line in read_jsonl_zst(self._graph_file): - if line.strip(): - node = FileNode.model_validate_json(line) - self._nodes[node.path] = node - 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 with self._io_lock: + if not self._graph_file.exists(): + return + try: + nodes, inverse, pending = await complete_in_thread( + self._load_sync, + self._graph_file, + self._nodes, + ) + self._nodes = nodes + self._inverse = inverse + self._pending = pending + 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: - try: - write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values())) - 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}") + async with self._io_lock: + try: + await complete_in_thread(self._dump_sync) + 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}") + + @classmethod + def _load_sync( + cls, + path: Path, + existing: dict[str, FileNode], + ) -> tuple[dict[str, FileNode], dict[str, set[str]], dict[str, set[str]]]: + """Restore nodes and rebuild adjacency entirely outside the event loop.""" + nodes = dict(existing) + for line in read_jsonl_zst(path): + if stripped := line.strip(): + node = FileNode.model_validate_json(stripped) + nodes[node.path] = node + + inverse: dict[str, set[str]] = {} + pending: dict[str, set[str]] = {} + for source, node in nodes.items(): + for target in cls._targets(node): + bucket = inverse if target in nodes else pending + bucket.setdefault(target, set()).add(source) + return nodes, inverse, pending + + def _dump_sync(self) -> None: + """Serialize and atomically publish the graph while its state is locked.""" + write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values())) # -- Internals --------------------------------------------------------- @@ -78,68 +111,89 @@ def _scope_match(self, target: str, scope: LinkScopeEnum | str) -> bool: # -- Node CRUD --------------------------------------------------------- async def upsert_nodes(self, nodes: list[FileNode]) -> None: - for node in nodes: - path = node.path - old = self._nodes.get(path) - if old is not None: - for target in self._targets(old): - self._remove_edge(path, target) - self._nodes[path] = node - for target in self._targets(node): - self._add_edge(path, target) - promoted = self._pending.pop(path, None) - if promoted: - self._inverse.setdefault(path, set()).update(promoted) + async with self._io_lock: + for node in nodes: + path = node.path + old = self._nodes.get(path) + if old is not None: + for target in self._targets(old): + self._remove_edge(path, target) + self._nodes[path] = node + for target in self._targets(node): + self._add_edge(path, target) + promoted = self._pending.pop(path, None) + if promoted: + self._inverse.setdefault(path, set()).update(promoted) async def delete_nodes(self, paths: list[str]) -> None: - for path in paths: - node = self._nodes.pop(path, None) - if node is None: - continue - for target in self._targets(node): - self._remove_edge(path, target) - demoted = self._inverse.pop(path, None) - if demoted: - self._pending.setdefault(path, set()).update(demoted) + async with self._io_lock: + for path in paths: + node = self._nodes.pop(path, None) + if node is None: + continue + for target in self._targets(node): + self._remove_edge(path, target) + demoted = self._inverse.pop(path, None) + if demoted: + self._pending.setdefault(path, set()).update(demoted) async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]: - if paths is None: - return list(self._nodes.values()) - return [self._nodes[p] for p in paths if p in self._nodes] + async with self._io_lock: + if paths is None: + return list(self._nodes.values()) + return [self._nodes[p] for p in paths if p in self._nodes] async def rebuild_links(self) -> None: - self._inverse.clear() - self._pending.clear() - for src, node in self._nodes.items(): - for target in self._targets(node): - self._add_edge(src, target) + async with self._io_lock: + nodes, inverse, pending = await complete_in_thread(self._rebuild_links_sync, self._nodes) + self._nodes = nodes + self._inverse = inverse + self._pending = pending + + @classmethod + def _rebuild_links_sync( + cls, + existing: dict[str, FileNode], + ) -> tuple[dict[str, FileNode], dict[str, set[str]], dict[str, set[str]]]: + """Rebuild adjacency from an in-memory node generation off-loop.""" + nodes = dict(existing) + inverse: dict[str, set[str]] = {} + pending: dict[str, set[str]] = {} + for source, node in nodes.items(): + for target in cls._targets(node): + bucket = inverse if target in nodes else pending + bucket.setdefault(target, set()).add(source) + return nodes, inverse, pending async def clear(self): - self._nodes.clear() - self._inverse.clear() - self._pending.clear() - self._graph_file.unlink(missing_ok=True) + async with self._io_lock: + self._nodes.clear() + self._inverse.clear() + self._pending.clear() + self._graph_file.unlink(missing_ok=True) # -- Link access ------------------------------------------------------- async def get_outlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]: - scope = self._normalize_scope(scope) - node = self._nodes.get(path) - if node is None: - return [] - return [lnk for lnk in node.links if lnk.target_path and self._scope_match(lnk.target_path, scope)] + async with self._io_lock: + scope = self._normalize_scope(scope) + node = self._nodes.get(path) + if node is None: + return [] + return [lnk for lnk in node.links if lnk.target_path and self._scope_match(lnk.target_path, scope)] async def get_inlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]: - scope = self._normalize_scope(scope) - sources: set[str] = set() - if scope in (LinkScopeEnum.REAL, LinkScopeEnum.ALL): - sources |= self._inverse.get(path, set()) - if scope in (LinkScopeEnum.VIRTUAL, LinkScopeEnum.ALL): - sources |= self._pending.get(path, set()) - return [ - link - for src in sorted(sources) - if src in self._nodes - for link in self._nodes[src].links - if link.target_path == path - ] + async with self._io_lock: + scope = self._normalize_scope(scope) + sources: set[str] = set() + if scope in (LinkScopeEnum.REAL, LinkScopeEnum.ALL): + sources |= self._inverse.get(path, set()) + if scope in (LinkScopeEnum.VIRTUAL, LinkScopeEnum.ALL): + sources |= self._pending.get(path, set()) + return [ + link + for src in sorted(sources) + if src in self._nodes + for link in self._nodes[src].links + if link.target_path == path + ] diff --git a/reme/components/file_graph/nx_file_graph.py b/reme/components/file_graph/nx_file_graph.py index 27257925..decbb4c9 100644 --- a/reme/components/file_graph/nx_file_graph.py +++ b/reme/components/file_graph/nx_file_graph.py @@ -1,12 +1,15 @@ """Networkx file-graph backend.""" +import asyncio import pickle from pathlib import Path +from uuid import uuid4 from .base_file_graph import BaseFileGraph from ..component_registry import R from ...enumeration import LinkScopeEnum from ...schema import FileLink, FileNode +from ...utils.async_utils import complete_in_thread @R.register("nx") @@ -23,30 +26,48 @@ def __init__(self, **kwargs): except ImportError as exc: raise ImportError("NxFileGraph requires networkx — pip install networkx") from exc self._graph = nx.MultiDiGraph() + self._io_lock = asyncio.Lock() self.component_metadata_path.mkdir(parents=True, exist_ok=True) self._graph_file: Path = self.component_metadata_path / f"{self.name}.pkl" # -- Lifecycle --------------------------------------------------------- async def load(self) -> None: - if not self._graph_file.exists(): - return - try: - with open(self._graph_file, "rb") as f: - self._graph = pickle.load(f) - 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 with self._io_lock: + if not self._graph_file.exists(): + return + try: + graph, real_count = await complete_in_thread(self._load_sync) + self._graph = graph + self.logger.info(f"Loaded {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: + async with self._io_lock: + try: + real_count = await complete_in_thread(self._dump_sync) + self.logger.info(f"Saved {real_count} nodes to {self._graph_file}") + except Exception as e: + self.logger.exception(f"Failed to write {self._graph_file}: {e}") + + def _load_sync(self): + """Load and count a NetworkX checkpoint outside the event loop.""" + with open(self._graph_file, "rb") as file: + graph = pickle.load(file) + real_count = sum(1 for _, data in graph.nodes(data=True) if "node" in data) + return graph, real_count + + def _dump_sync(self) -> int: + """Serialize and atomically publish the locked graph off-loop.""" + tmp = self._graph_file.with_name(f".{self._graph_file.name}.{uuid4().hex}.tmp") try: - tmp = self._graph_file.with_suffix(".tmp") - with open(tmp, "wb") as f: - pickle.dump(self._graph, f, protocol=pickle.HIGHEST_PROTOCOL) + with open(tmp, "wb") as file: + pickle.dump(self._graph, file, protocol=pickle.HIGHEST_PROTOCOL) tmp.replace(self._graph_file) - 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}") + return self._real_count() + finally: + tmp.unlink(missing_ok=True) # -- Internals --------------------------------------------------------- @@ -69,29 +90,37 @@ def _scope_match(self, key: str, scope: LinkScopeEnum) -> bool: # -- Node CRUD --------------------------------------------------------- async def upsert_nodes(self, nodes: list[FileNode]) -> None: - for node in nodes: - path = node.path - if self._graph.has_node(path): - self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True))) - self._graph.add_node(path, node=node) # promotes virtual placeholder - self._graph.add_edges_from(self._edges_from(path, node)) + async with self._io_lock: + for node in nodes: + path = node.path + if self._graph.has_node(path): + self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True))) + self._graph.add_node(path, node=node) # promotes virtual placeholder + self._graph.add_edges_from(self._edges_from(path, node)) async def delete_nodes(self, paths: list[str]) -> None: - for path in paths: - if not self._graph.has_node(path): - continue - 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: - self._graph.remove_node(path) + async with self._io_lock: + for path in paths: + if not self._graph.has_node(path): + continue + 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: + self._graph.remove_node(path) async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]: - view = self._graph.nodes - if paths is None: - return [d["node"] for _, d in view(data=True) if "node" in d] - return [view[p]["node"] for p in paths if p in view and "node" in view[p]] + async with self._io_lock: + view = self._graph.nodes + if paths is None: + return [d["node"] for _, d in view(data=True) if "node" in d] + return [view[p]["node"] for p in paths if p in view and "node" in view[p]] async def rebuild_links(self) -> None: + async with self._io_lock: + await complete_in_thread(self._rebuild_links_sync) + + def _rebuild_links_sync(self) -> None: + """Rebuild NetworkX edges outside the event-loop thread.""" self._graph.remove_edges_from(list(self._graph.edges(keys=True))) virtual = [n for n, d in self._graph.nodes(data=True) if "node" not in d] self._graph.remove_nodes_from(virtual) @@ -99,23 +128,26 @@ async def rebuild_links(self) -> None: self._graph.add_edges_from(self._edges_from(path, data["node"])) async def clear(self): - self._graph.clear() - self._graph_file.unlink(missing_ok=True) + async with self._io_lock: + self._graph.clear() + self._graph_file.unlink(missing_ok=True) # -- Link access ------------------------------------------------------- async def get_outlinks(self, path: str, scope: LinkScopeEnum = LinkScopeEnum.REAL) -> list[FileLink]: - view = self._graph.nodes - if path not in view or "node" not in view[path]: - return [] - return [ - d["link"] - for _, tgt, d in self._graph.out_edges(path, data=True) - if "link" in d and self._scope_match(tgt, scope) - ] + async with self._io_lock: + view = self._graph.nodes + if path not in view or "node" not in view[path]: + return [] + return [ + d["link"] + for _, tgt, d in self._graph.out_edges(path, data=True) + if "link" in d and self._scope_match(tgt, scope) + ] async def get_inlinks(self, path: str, scope: LinkScopeEnum = LinkScopeEnum.REAL) -> list[FileLink]: - view = self._graph.nodes - if path not in view or not self._scope_match(path, scope): - return [] - return [d["link"] for _, _, d in self._graph.in_edges(path, data=True) if "link" in d] + async with self._io_lock: + view = self._graph.nodes + if path not in view or not self._scope_match(path, scope): + return [] + return [d["link"] for _, _, d in self._graph.in_edges(path, data=True) if "link" in d] diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 12684e81..e07d5726 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -6,13 +6,13 @@ from contextlib import suppress from uuid import uuid4 -import aiofiles import numpy as np from .base_file_store import BaseFileStore from .local_file_store import LocalFileStore from ..component_registry import R from ...schema import FileChunk, FileNode +from ...utils.async_utils import complete_in_thread @R.register("faiss") @@ -410,6 +410,7 @@ def _chunks_embedding_digest(self) -> str: digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) return digest.hexdigest() + @BaseFileStore.serialized async def load(self) -> None: """Load chunks via the parent, then attach FAISS state (sidecar or rebuild).""" await super().load() @@ -417,9 +418,13 @@ async def load(self) -> None: self._faiss_index = None return if not await self._try_load_sidecar(): - self._rebuild_index() + await complete_in_thread(self._rebuild_index) async def _try_load_sidecar(self) -> bool: + """Load and validate the complete FAISS sidecar outside the event loop.""" + return await complete_in_thread(self._try_load_sidecar_sync) + + def _try_load_sidecar_sync(self) -> bool: """Read the binary index plus id-map sidecar. On any mismatch or read error, wipe the partial files so the caller can rebuild from chunks cleanly. @@ -459,8 +464,7 @@ async def _try_load_sidecar(self) -> bool: raise ValueError( f"FAISS HNSW M mismatch: persisted={persisted_m}, configured={self.hnsw_m}", ) - async with aiofiles.open(self.faiss_idmap_path, encoding=self.encoding) as f: - data = json.loads(await f.read()) + data = json.loads(self.faiss_idmap_path.read_text(encoding=self.encoding)) id_map = list(data.get("id_map", [])) if len(id_map) != index.ntotal: raise ValueError(f"id_map size {len(id_map)} != index ntotal {index.ntotal}") @@ -506,6 +510,7 @@ async def _try_load_sidecar(self) -> bool: self.faiss_idmap_path.unlink(missing_ok=True) return False + @BaseFileStore.serialized async def _dump_owned_state(self) -> None: """Persist chunks and the FAISS sidecar, excluding dependency snapshots.""" async with self._faiss_dump_lock: @@ -513,7 +518,7 @@ async def _dump_owned_state(self) -> None: if self._faiss_index is None or self.embedding_store is None: return try: - self._compact_if_needed() + await complete_in_thread(self._compact_if_needed) await self._write_sidecar() self.logger.info(f"Saved FAISS index: {self._faiss_index.ntotal} vectors to {self.faiss_path}") except Exception as e: @@ -521,6 +526,10 @@ async def _dump_owned_state(self) -> None: raise async def _write_sidecar(self) -> None: + """Persist both FAISS sidecar files outside the event loop.""" + await complete_in_thread(self._write_sidecar_sync) + + def _write_sidecar_sync(self) -> None: token = uuid4().hex tmp_index = self.faiss_path.with_name(f".{self.faiss_path.name}.{token}.tmp") tmp_idmap = self.faiss_idmap_path.with_name(f".{self.faiss_idmap_path.name}.{token}.tmp") @@ -536,8 +545,7 @@ async def _write_sidecar(self) -> None: ) try: self._faiss.write_index(self._faiss_index, str(tmp_index)) - async with aiofiles.open(tmp_idmap, "w", encoding=self.encoding) as f: - await f.write(payload) + tmp_idmap.write_text(payload, encoding=self.encoding) # Publish only after both parts of the sidecar have been written successfully. tmp_index.replace(self.faiss_path) @@ -607,7 +615,7 @@ async def delete(self, path: str | list[str]) -> None: deleted_ids = [cid for n in nodes for cid in n.chunk_ids] await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes if nodes: - self._mutation_generation += 1 + self._advance_mutation_generation() if self._embedding_rebuild_pending or self._faiss_index is None: return for cid in deleted_ids: diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index cab8024f..3707a396 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -73,6 +73,7 @@ def __init__( self._embedding_rebuild_pending = bool(embedding_rebuild_required) self._embedding_space_generation = 0 self._mutation_generation = 0 + self._checkpoint_generation = 0 self._closing = False # -- lifecycle ------------------------------------------------------------ @@ -213,17 +214,20 @@ async def _get_query_embedding(self, query: str) -> np.ndarray | None: # -- persistence ---------------------------------------------------------- + @BaseFileStore.serialized async def load(self) -> None: """Load chunks from the JSONL file into memory; missing file is a no-op.""" started_at = time.monotonic() chunk_load_started_at = time.monotonic() if self.chunks_path.exists(): try: - for line in read_jsonl_zst(self.chunks_path, self.encoding): - line = line.strip() - if line: - chunk = self._deserialize_chunk(line) - self.file_chunks[chunk.id] = chunk + chunks = await complete_in_thread( + self._load_chunks_sync, + self.chunks_path, + self.encoding, + self.file_chunks, + ) + self.file_chunks = chunks self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}") except Exception as e: self.logger.exception(f"Failed to load {self.chunks_path}: {e}") @@ -291,6 +295,16 @@ def _deserialize_chunk(line: str) -> FileChunk: payload["embedding"] = np.frombuffer(raw, dtype=_EMBEDDING_F16_DTYPE) return FileChunk.model_validate(payload) + @classmethod + def _load_chunks_sync(cls, path, encoding: str, existing: dict[str, FileChunk]) -> dict[str, FileChunk]: + """Read, decompress, and parse a complete chunk checkpoint off-loop.""" + chunks = dict(existing) + for line in read_jsonl_zst(path, encoding): + if stripped := line.strip(): + chunk = cls._deserialize_chunk(stripped) + chunks[chunk.id] = chunk + return chunks + @staticmethod def _serialize_chunk(chunk: FileChunk) -> str: """Serialize embeddings without expanding float16 values into Python floats.""" @@ -509,6 +523,10 @@ async def _backfill_missing_embeddings_inner(self, *, skip_health_check: bool, s batch = missing[start : start + batch_size] await self.embedding_store.get_node_embeddings(batch) self._drop_stale_embeddings(batch, "backfill") + # A checkpoint snapshot may be copying chunks in a worker while + # this background task applies embeddings on the event loop. + # Advance the generation so a mixed snapshot is discarded. + self._checkpoint_generation += 1 processed += len(batch) batch_count += 1 next_percent = self._log_progress("embedding backfill", processed, total, next_percent) @@ -617,19 +635,28 @@ async def _rebuild_keyword_index(self, docs: dict[str, str]) -> None: elapsed = time.monotonic() - started_at self.logger.info(f"{self.name}: keyword index rebuild complete: total={total}, elapsed={elapsed:.2f}s") + @BaseFileStore.serialized async def _dump_owned_state(self) -> None: """Persist state owned by this store, excluding dependency snapshots.""" try: - # Keep snapshotting synchronous so concurrent mutation cannot produce a - # mixed-generation checkpoint. Move it off-loop only if profiling shows - # this copy, rather than serialization/compression, is a material stall. - chunks = tuple(chunk.model_copy(deep=True) for chunk in self.file_chunks.values()) + # The maintenance guard excludes regular mutations. Embedding + # backfill intentionally does network I/O outside that guard, so + # retry if it publishes a new generation during the worker copy. + while True: + generation = self._checkpoint_generation + chunks = await complete_in_thread(self._snapshot_chunks_sync) + if generation == self._checkpoint_generation: + break await complete_in_thread(self._dump_chunks_sync, chunks) 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}") raise + def _snapshot_chunks_sync(self) -> tuple[FileChunk, ...]: + """Deep-copy the live chunk generation outside the event-loop thread.""" + return tuple(chunk.model_copy(deep=True) for chunk in self.file_chunks.values()) + def _dump_chunks_sync(self, chunks: tuple[FileChunk, ...]) -> None: write_jsonl_zst( self.chunks_path, @@ -676,7 +703,12 @@ async def upsert(self, files: list[tuple[FileNode, list[FileChunk]]]) -> None: await self.keyword_index.delete_docs(list(old_chunk_ids)) if self.keyword_index and keyword_docs: await self.keyword_index.add_docs(keyword_docs) + self._advance_mutation_generation() + + def _advance_mutation_generation(self) -> None: + """Mark a source mutation for rebuild retries and checkpoint snapshots.""" self._mutation_generation += 1 + self._checkpoint_generation += 1 def _stage_upsert( self, @@ -767,7 +799,7 @@ async def delete(self, path: str | list[str]) -> None: nodes: list[FileNode] = await self.file_graph.get_nodes(paths) await self._delete_nodes(nodes) if nodes: - self._mutation_generation += 1 + self._advance_mutation_generation() async def _delete_nodes(self, nodes: list[FileNode]) -> None: """Delete already-resolved nodes and their chunks. @@ -814,7 +846,7 @@ async def clear(self) -> None: if self.keyword_index: await self.keyword_index.clear() await self.file_graph.clear() - self._mutation_generation += 1 + self._advance_mutation_generation() # -- search --------------------------------------------------------------- diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index afad4b4d..88945a93 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -6,7 +6,6 @@ import time from uuid import uuid4 -import aiofiles import numpy as np from .base_file_store import BaseFileStore @@ -246,6 +245,7 @@ def _chunks_embedding_digest(self) -> str: digest.update(np.asarray(self.file_chunks[cid].embedding, dtype=np.float16).tobytes()) return digest.hexdigest() + @BaseFileStore.serialized async def load(self) -> None: """Load chunks via the parent, then attach the zvec collection (open or rebuild).""" await super().load() @@ -253,9 +253,13 @@ async def load(self) -> None: self._collection = None return if not await self._try_open_collection(): - self._rebuild_collection() + await complete_in_thread(self._rebuild_collection) async def _try_open_collection(self) -> bool: + """Open and validate the complete zvec checkpoint off-loop.""" + return await complete_in_thread(self._try_open_collection_sync) + + def _try_open_collection_sync(self) -> bool: """Open the persisted collection and validate it against the chunks. On any mismatch or open error the collection directory and sidecar are @@ -270,8 +274,7 @@ async def _try_open_collection(self) -> bool: return False collection = None try: - async with aiofiles.open(self.zvec_sidecar_path, encoding=self.encoding) as f: - sidecar = json.loads(await f.read()) + sidecar = json.loads(self.zvec_sidecar_path.read_text(encoding=self.encoding)) if sidecar.get("digest") != self._chunks_embedding_digest(): raise ValueError("zvec sidecar embedding digest does not match persisted chunks") indexed_ids = set(sidecar.get("ids", [])) @@ -333,13 +336,14 @@ def _verify_collection_contents(self, collection, expected_ids: set[str]) -> Non f"elapsed={time.monotonic() - started_at:.3f}s", ) + @BaseFileStore.serialized 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: - self._collection.flush() + await complete_in_thread(self._collection.flush) await self._write_sidecar() self.logger.info(f"Saved zvec collection: {len(self._indexed_ids)} vectors to {self.zvec_path}") except Exception as e: @@ -348,6 +352,10 @@ async def _dump_owned_state(self) -> None: async def _write_sidecar(self) -> None: """Atomically write the digest sidecar binding the collection to the chunk generation.""" + await complete_in_thread(self._write_sidecar_sync) + + def _write_sidecar_sync(self) -> None: + """Build and atomically publish the zvec sidecar off-loop.""" tmp = self.zvec_sidecar_path.with_name(f".{self.zvec_sidecar_path.name}.{uuid4().hex}.tmp") payload = json.dumps( { @@ -356,8 +364,7 @@ async def _write_sidecar(self) -> None: }, ) try: - async with aiofiles.open(tmp, "w", encoding=self.encoding) as f: - await f.write(payload) + tmp.write_text(payload, encoding=self.encoding) tmp.replace(self.zvec_sidecar_path) finally: tmp.unlink(missing_ok=True) @@ -416,7 +423,7 @@ async def delete(self, path: str | list[str]) -> None: deleted_ids = [cid for n in nodes for cid in n.chunk_ids] await self._delete_nodes(nodes) # reuse resolved nodes; avoids a second get_nodes if nodes: - self._mutation_generation += 1 + self._advance_mutation_generation() if self._embedding_rebuild_pending: return self._delete_docs(deleted_ids) diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 0169726f..03ef4e4d 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -56,6 +56,7 @@ def __init__(self, k1: float = 1.5, b: float = 0.75, index_version: str = "v1", # IDF cache; invalidated whenever live-doc count or postings change. self._idf_cache: dict[int, float] = {} self._dump_lock = asyncio.Lock() + self._state_lock = asyncio.Lock() # -- Properties ----------------------------------------------------------- @@ -249,6 +250,12 @@ async def add_docs(self, docs_dict: dict[str, str]) -> None: if not docs_dict: return + async with self._state_lock: + self._add_docs_sync(docs_dict) + + def _add_docs_sync(self, docs_dict: dict[str, str]) -> None: + """Mutate the inverted index while the caller owns the state lock.""" + new_doc_ids: list[str] = [] new_doc_lens: list[int] = [] new_doc_token_ids: list[np.ndarray] = [] @@ -276,9 +283,10 @@ async def add_docs(self, docs_dict: dict[str, str]) -> None: async def delete_docs(self, doc_ids: list[str]) -> None: """Lazy-delete a batch of doc_ids; physical reclaim happens in optimize_index.""" - for doc_id in doc_ids: - self._remove_doc(doc_id) - self._idf_cache = {} + async with self._state_lock: + for doc_id in doc_ids: + self._remove_doc(doc_id) + self._idf_cache = {} def _score_query(self, query_ids: list[int], n_docs: int) -> np.ndarray: """Compute BM25 scores across all docs; deleted docs zeroed out.""" @@ -361,14 +369,12 @@ def _restore(self, data: dict) -> None: async def dump(self) -> None: """Persist the index via temp file + atomic rename to avoid torn writes.""" async with self._dump_lock: - if self.n_docs == 0 and not self.vocab: - self.index_file.unlink(missing_ok=True) - return try: - # Keep snapshotting synchronous so the worker receives one coherent - # index generation. Move it off-loop only if profiling identifies this - # copy, rather than pickle/file I/O, as a material event-loop stall. - snapshot = self._snapshot() + async with self._state_lock: + if self.n_docs == 0 and not self.vocab: + self.index_file.unlink(missing_ok=True) + return + snapshot = await complete_in_thread(self._snapshot) await complete_in_thread(self._dump_sync, snapshot) self.logger.info(f"Saved {self.n_docs} docs to {self.index_file}") except Exception as e: @@ -389,14 +395,17 @@ async def load(self) -> None: """Load from disk; missing file is a no-op, corrupt file resets state.""" if not self.index_file.exists(): return - try: - data = await asyncio.to_thread(self._load_sync) - self._restore(data) - 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}") - self.index_file.unlink(missing_ok=True) - await self.clear() + async with self._dump_lock: + try: + data = await asyncio.to_thread(self._load_sync) + async with self._state_lock: + self._restore(data) + 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}") + self.index_file.unlink(missing_ok=True) + async with self._state_lock: + self._clear_state() def _load_sync(self) -> dict: with open(self.index_file, "rb") as file: @@ -405,16 +414,21 @@ def _load_sync(self) -> dict: async def clear(self) -> None: """Reset in-memory state and remove the persisted file.""" async with self._dump_lock: - self.vocab = {} - self._doc_ids = [] - self._doc_id_to_idx = {} - self._doc_lens = np.zeros(0, dtype=np.int32) - self._deleted = np.zeros(0, dtype=bool) - self._doc_token_ids = [] - self._posting_doc_idxs = {} - self._posting_tfs = {} - self._idf_cache = {} - self.index_file.unlink(missing_ok=True) + async with self._state_lock: + self._clear_state() + self.index_file.unlink(missing_ok=True) + + def _clear_state(self) -> None: + """Reset every in-memory field while the caller owns the state lock.""" + self.vocab = {} + self._doc_ids = [] + self._doc_id_to_idx = {} + self._doc_lens = np.zeros(0, dtype=np.int32) + self._deleted = np.zeros(0, dtype=bool) + self._doc_token_ids = [] + self._posting_doc_idxs = {} + self._posting_tfs = {} + self._idf_cache = {} # -- Compaction ----------------------------------------------------------- @@ -475,28 +489,30 @@ def _compact_docs( async def optimize_index(self) -> None: """Physically reclaim deleted docs and unused vocab entries.""" - if self._deleted.size == 0: - return - active_mask = ~self._deleted - if not active_mask.any(): - await self.clear() - return + async with self._state_lock: + if self._deleted.size == 0: + return + active_mask = ~self._deleted + if not active_mask.any(): + self._clear_state() + self.index_file.unlink(missing_ok=True) + return - old_to_new_idx, n_active = self._build_idx_remap(active_mask) - new_vocab, old_tid_to_new = self._compact_vocab(active_mask) - new_posting_idxs, new_posting_tfs = self._compact_postings( - active_mask, - old_to_new_idx, - old_tid_to_new, - ) - new_doc_ids, new_doc_token_ids = self._compact_docs(active_mask, old_tid_to_new) - - self.vocab = new_vocab - self._doc_ids = new_doc_ids - self._doc_id_to_idx = {doc_id: i for i, doc_id in enumerate(new_doc_ids)} - self._doc_lens = self._doc_lens[active_mask].astype(np.int32, copy=True) - self._deleted = np.zeros(n_active, dtype=bool) - self._doc_token_ids = new_doc_token_ids - self._posting_doc_idxs = new_posting_idxs - self._posting_tfs = new_posting_tfs - self._idf_cache = {} + old_to_new_idx, n_active = self._build_idx_remap(active_mask) + new_vocab, old_tid_to_new = self._compact_vocab(active_mask) + new_posting_idxs, new_posting_tfs = self._compact_postings( + active_mask, + old_to_new_idx, + old_tid_to_new, + ) + new_doc_ids, new_doc_token_ids = self._compact_docs(active_mask, old_tid_to_new) + + self.vocab = new_vocab + self._doc_ids = new_doc_ids + self._doc_id_to_idx = {doc_id: i for i, doc_id in enumerate(new_doc_ids)} + self._doc_lens = self._doc_lens[active_mask].astype(np.int32, copy=True) + self._deleted = np.zeros(n_active, dtype=bool) + self._doc_token_ids = new_doc_token_ids + self._posting_doc_idxs = new_posting_idxs + self._posting_tfs = new_posting_tfs + self._idf_cache = {} diff --git a/tests/unit/test_embedded_consumer_compat.py b/tests/unit/test_embedded_consumer_compat.py index 40d12795..5c3edecf 100644 --- a/tests/unit/test_embedded_consumer_compat.py +++ b/tests/unit/test_embedded_consumer_compat.py @@ -3,14 +3,17 @@ # pylint: disable=protected-access import asyncio +import threading import pytest from reme import ReMe from reme.components.agent_wrapper import AsAgentWrapper from reme.components.as_llm import BaseAsLLM, DashScopeAsLLM +from reme.components.file_store import local_file_store as local_file_store_module from reme.enumeration import ComponentEnum from reme.schema import FileNode +from reme.utils.jsonl_zst import write_jsonl_zst def _qwenpaw_style_config(workspace_dir: str) -> dict: @@ -59,6 +62,85 @@ def _file_graph_config(workspace_dir: str) -> dict: } +def _file_store_config(workspace_dir: str) -> dict: + """Return the persistent component graph used by an embedded local store.""" + return { + "workspace_dir": workspace_dir, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "http"}, + "components": { + "tokenizer": {"default": {"backend": "regex"}}, + "keyword_index": {"default": {"backend": "bm25", "tokenizer": "default"}}, + "file_graph": {"default": {"backend": "local"}}, + "file_store": { + "default": { + "backend": "local", + "embedding_store": "", + "keyword_index": "default", + "file_graph": "default", + }, + }, + }, + } + + +def test_embedded_reme_lifecycle_keeps_checkpoint_io_off_loop(tmp_path, monkeypatch): + """Real ReMe start and close await worker-backed checkpoint operations.""" + app = ReMe(**_file_store_config(str(tmp_path))) + store = app.context.components[ComponentEnum.FILE_STORE]["default"] + graph = app.context.components[ComponentEnum.FILE_GRAPH]["default"] + store.chunks_path.parent.mkdir(parents=True, exist_ok=True) + write_jsonl_zst(store.chunks_path, []) + + async def exercise_api() -> None: + loop_thread = threading.get_ident() + load_entered = threading.Event() + load_release = threading.Event() + load_threads = [] + original_reader = local_file_store_module.read_jsonl_zst + + def blocking_reader(*args, **kwargs): + load_threads.append(threading.get_ident()) + load_entered.set() + assert load_release.wait(timeout=2) + yield from original_reader(*args, **kwargs) + + monkeypatch.setattr(local_file_store_module, "read_jsonl_zst", blocking_reader) + start_task = asyncio.create_task(app.start()) + assert await asyncio.to_thread(load_entered.wait, 1) + assert load_threads and load_threads[0] != loop_thread + for _ in range(5): + await asyncio.sleep(0) + assert not start_task.done() + load_release.set() + await start_task + + dump_entered = threading.Event() + dump_release = threading.Event() + dump_threads = [] + original_dump = graph._dump_sync + + def blocking_dump(): + dump_threads.append(threading.get_ident()) + dump_entered.set() + assert dump_release.wait(timeout=2) + return original_dump() + + monkeypatch.setattr(graph, "_dump_sync", blocking_dump) + close_task = asyncio.create_task(app.close()) + assert await asyncio.to_thread(dump_entered.wait, 1) + assert dump_threads and dump_threads[0] != loop_thread + for _ in range(5): + await asyncio.sleep(0) + assert not close_task.done() + dump_release.set() + await close_task + + asyncio.run(exercise_api()) + + def test_qwenpaw_style_config_preserves_optional_defaults(tmp_path): """New application fields remain optional for existing embedded configs.""" app = ReMe(**_qwenpaw_style_config(str(tmp_path))) diff --git a/tests/unit/test_file_catalog.py b/tests/unit/test_file_catalog.py index 06c1c90a..46a6054d 100644 --- a/tests/unit/test_file_catalog.py +++ b/tests/unit/test_file_catalog.py @@ -5,6 +5,7 @@ import asyncio import os import tempfile +import threading import pytest @@ -174,6 +175,42 @@ async def run(): asyncio.run(run()) +def test_catalog_lifecycle_checkpoint_work_runs_off_loop(tmp_path, monkeypatch): + """Catalog start and close execute full compressed checkpoints in workers.""" + + async def run(): + with temp_chdir(tmp_path): + seed = LocalFileCatalog() + await seed.start() + await seed.upsert([make_node("a.md")]) + await seed.close() + + catalog = LocalFileCatalog() + loop_thread = threading.get_ident() + load_threads = [] + dump_threads = [] + original_load = catalog._read_jsonl_sync + original_dump = catalog._write_jsonl_sync + + def observed_load(*args): + load_threads.append(threading.get_ident()) + return original_load(*args) + + def observed_dump(): + dump_threads.append(threading.get_ident()) + return original_dump() + + monkeypatch.setattr(catalog, "_read_jsonl_sync", observed_load) + monkeypatch.setattr(catalog, "_write_jsonl_sync", observed_dump) + await catalog.start() + await catalog.close() + + assert load_threads and all(thread_id != loop_thread for thread_id in load_threads) + assert dump_threads and all(thread_id != loop_thread for thread_id in dump_threads) + + asyncio.run(run()) + + if __name__ == "__main__": print("\n=== FileCatalog Tests ===") for backend in BACKENDS: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 26f2dce6..fb2e79f2 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -1291,6 +1291,125 @@ def blocking_dump(_chunks): run(go()) +def test_start_and_close_checkpoint_io_keep_event_loop_responsive(monkeypatch): + """Public lifecycle methods keep synchronous checkpoint work off-loop.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_nonblocking_lifecycle") + store.chunks_path.parent.mkdir(parents=True, exist_ok=True) + write_jsonl_zst(store.chunks_path, []) + loop_thread = threading.get_ident() + + load_entered = threading.Event() + load_release = threading.Event() + load_threads = [] + original_reader = local_file_store_module.read_jsonl_zst + + def blocking_reader(*args, **kwargs): + load_threads.append(threading.get_ident()) + load_entered.set() + assert load_release.wait(timeout=2) + yield from original_reader(*args, **kwargs) + + monkeypatch.setattr(local_file_store_module, "read_jsonl_zst", blocking_reader) + start_task = asyncio.create_task(store.start()) + assert await asyncio.to_thread(load_entered.wait, 1) + assert load_threads == [load_threads[0]] + assert load_threads[0] != loop_thread + for _ in range(5): + await asyncio.sleep(0) + assert not start_task.done() + load_release.set() + await start_task + + graph = store.file_graph + dump_entered = threading.Event() + dump_release = threading.Event() + dump_threads = [] + original_graph_dump = graph._dump_sync + + def blocking_graph_dump(): + dump_threads.append(threading.get_ident()) + dump_entered.set() + assert dump_release.wait(timeout=2) + return original_graph_dump() + + monkeypatch.setattr(graph, "_dump_sync", blocking_graph_dump) + close_task = asyncio.create_task(store.close()) + assert await asyncio.to_thread(dump_entered.wait, 1) + assert dump_threads == [dump_threads[0]] + assert dump_threads[0] != loop_thread + for _ in range(5): + await asyncio.sleep(0) + assert not close_task.done() + dump_release.set() + await close_task + + run(go()) + + +def test_checkpoint_snapshot_is_built_off_event_loop(monkeypatch): + """Deep-copying a chunk generation runs in the worker before compression.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_nonblocking_snapshot") + await store.start() + store.file_chunks["a"] = chunk("a", "a.md", "alpha") + loop_thread = threading.get_ident() + snapshot_threads = [] + original_snapshot = store._snapshot_chunks_sync + + def observed_snapshot(): + snapshot_threads.append(threading.get_ident()) + return original_snapshot() + + monkeypatch.setattr(store, "_snapshot_chunks_sync", observed_snapshot) + await store._dump_owned_state() + + assert snapshot_threads == [snapshot_threads[0]] + assert snapshot_threads[0] != loop_thread + await store.close() + + run(go()) + + +def test_checkpoint_snapshot_retries_concurrent_embedding_generation(monkeypatch): + """A backfill publication during worker copy cannot produce a mixed checkpoint.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = _new_local_store("t_snapshot_generation_retry") + await store.start() + store.file_chunks["a"] = chunk("a", "a.md", "alpha") + first_entered = threading.Event() + release_first = threading.Event() + snapshot_calls = 0 + original_snapshot = store._snapshot_chunks_sync + + def blocking_snapshot(): + nonlocal snapshot_calls + snapshot_calls += 1 + snapshot = original_snapshot() + if snapshot_calls == 1: + first_entered.set() + assert release_first.wait(timeout=2) + return snapshot + + monkeypatch.setattr(store, "_snapshot_chunks_sync", blocking_snapshot) + dump_task = asyncio.create_task(store._dump_owned_state()) + assert await asyncio.to_thread(first_entered.wait, 1) + store._checkpoint_generation += 1 + release_first.set() + await dump_task + + assert snapshot_calls == 2 + await store.close() + + run(go()) + + def test_checkpoint_dump_finishes_before_propagating_cancellation(monkeypatch): """Cancellation cannot leave an old checkpoint writer running in the background.""" diff --git a/tests/unit/test_keyword_index.py b/tests/unit/test_keyword_index.py index 061c97b7..c251b792 100644 --- a/tests/unit/test_keyword_index.py +++ b/tests/unit/test_keyword_index.py @@ -882,6 +882,32 @@ def blocking_dump_sync(snapshot): run(go()) +def test_dump_builds_snapshot_off_event_loop(): + """Large BM25 snapshot copies do not execute on the request event loop.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + bm25 = await create_bm25() + await bm25.add_docs({"d1": "alpha beta"}) + loop_thread = threading.get_ident() + snapshot_threads = [] + original_snapshot = bm25._snapshot + + def observed_snapshot(): + snapshot_threads.append(threading.get_ident()) + return original_snapshot() + + bm25._snapshot = observed_snapshot + await bm25.dump() + + assert snapshot_threads == [snapshot_threads[0]] + assert snapshot_threads[0] != loop_thread + bm25._snapshot = original_snapshot + await bm25.close() + + run(go()) + + # --------------------------------------------------------------------------- # # clear / optimize / reset_index # # --------------------------------------------------------------------------- #