mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 20:14:32 +00:00
fix(graphrag): take a source's write lock per chunk, and keep zero weights
Two builds of one source can overlap: the extraction lease is keyed by the
source's updated_at, and enabling a graph updates the source before it
dispatches, so a rebuild started while the last build runs gets a new key and
a lease of its own. Both builds could then pass a chunk's "done" check before
either committed and apply it twice -- doc_freq bumped twice, reproduced with
two live writers. A reset could also land in the middle of a chunk.
apply_chunk and delete_by_source now take a transaction-scoped advisory lock
keyed by the source before touching a row, as the schema bootstrap already
does for DDL. A single build's writes were already serial, so it loses
nothing; overlapping builds take turns chunk by chunk, and the second sees the
first's "done" row and returns (0, 0).
apply_chunk also still defaulted with `rel.get("weight") or 1.0`, turning an
explicit zero into a full-strength edge -- the conversion 3f774d81 removed from
add_edge and the ranker but missed here. Only a missing weight defaults now.
This commit is contained in:
1 parent
b3d7e63ae3
commit
ecf02d0d60
2 files changed
+119
-3
No files matched your search
@@ -108,6 +108,22 @@ def _is_connection_lost(exc: BaseException) -> bool:
|
||||
return isinstance(exc, (psycopg.OperationalError, psycopg.InterfaceError))
|
||||
|
||||
|
||||
def _lock_source(cursor, source_id: str) -> None:
|
||||
"""Serialize graph writes for one source until this transaction ends.
|
||||
|
||||
Writes within one build are already serial, but two builds of the same
|
||||
source can overlap: a rebuild dispatched while the last one is still
|
||||
running gets a new idempotency key, so its lease does not stop it. Without
|
||||
this, both could pass a chunk's "done" check before either commits and
|
||||
apply it twice. A transaction-scoped advisory lock keyed by the source
|
||||
makes them take turns chunk by chunk; the lock is released on commit or
|
||||
rollback, and a hash collision only makes two sources take turns.
|
||||
"""
|
||||
cursor.execute(
|
||||
"SELECT pg_advisory_xact_lock(hashtext(%s));", (f"graphrag:source:{source_id}",)
|
||||
)
|
||||
|
||||
|
||||
def _safe_rollback(conn) -> None:
|
||||
"""Roll back, tolerating a connection too broken to roll back."""
|
||||
try:
|
||||
@@ -672,7 +688,10 @@ class GraphStore:
|
||||
# committed, and the retry then replays this write: doc_freq
|
||||
# would be bumped twice and a second logical edge inserted
|
||||
# (graph_edges has no uniqueness constraint). The progress row
|
||||
# below is written in this transaction, so a replay sees it.
|
||||
# below is written in this transaction, so a replay sees it —
|
||||
# and so does an overlapping build, once the source lock makes
|
||||
# it wait for this one to commit.
|
||||
_lock_source(cursor, source_id)
|
||||
cursor.execute(
|
||||
"SELECT status FROM graph_ingest_progress "
|
||||
"WHERE source_id = %s AND chunk_id = %s;",
|
||||
@@ -713,7 +732,9 @@ class GraphStore:
|
||||
dst_id,
|
||||
type=rel.get("type"),
|
||||
description=rel.get("description"),
|
||||
weight=float(rel.get("weight") or 1.0),
|
||||
# Only a missing weight defaults: 0 is a real one,
|
||||
# and the ranker drops non-positive edges.
|
||||
weight=1.0 if rel.get("weight") is None else float(rel["weight"]),
|
||||
source_chunk_ids=[chunk_id],
|
||||
fact_embedding=rel.get("fact_embedding"),
|
||||
)
|
||||
@@ -1527,6 +1548,9 @@ class GraphStore:
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
# A reset while a build is still writing must not land in the
|
||||
# middle of one of its chunks.
|
||||
_lock_source(cursor, source_id)
|
||||
for table in (
|
||||
"graph_node_chunks",
|
||||
"graph_edges",
|
||||
|
||||
Reference in new issue
Block a user