Skip to content

Commit c3c238d

Browse files
remote-controlled-manremote-controlled-man
committed
fix(index): drain deferred relation vector shards
Consume the vector sync batch result instead of discarding it, drain deferred shards with a no-progress guard, and expose failed/deferred entity IDs on the relation-resolution result. Failed entities remain marked by the vector backend for its retry/readiness path. Signed-off-by: Steven Huang <[email protected]>
1 parent 249a2e0 commit c3c238d

3 files changed

Lines changed: 146 additions & 13 deletions

File tree

‎src/basic_memory/indexing/relation_resolution.py‎

Lines changed: 57 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -553,6 +553,11 @@ class ResolveRelationsResult:
553553
# Source notes whose links this run resolved; see
554554
# services.read_cache_changes.resolved_relations_read_cache_change.
555555
affected_entity_ids: frozenset[int] = frozenset()
556+
# Entities the vector backend could not finish. Failed entities remain marked
557+
# for the backend's retry/readiness path; the result exposes both sets instead
558+
# of silently reporting the relation refresh as fully synced.
559+
vector_sync_failed_entity_ids: frozenset[int] = frozenset()
560+
vector_sync_deferred_entity_ids: frozenset[int] = frozenset()
556561

557562
@property
558563
def resolved(self) -> int:
@@ -577,9 +582,11 @@ async def resolve_project_relations(
577582
genuine forward references and remain unresolved until their target note
578583
exists.
579584
580-
When ``vector_sync`` is set, callers that embed run it once after every
581-
pass settles: resolution rewrites the affected sources' relation search
582-
rows, and only changed chunks are re-embedded.
585+
When ``vector_sync`` is set, callers that embed run it after every pass
586+
settles and then drain any deferred shards. Resolution rewrites the affected
587+
sources' relation search rows, and only changed chunks are re-embedded.
588+
Failed entities remain marked by the backend's retry/readiness state and are
589+
exposed on the result.
583590
"""
584591
unresolved_before = await runtime.count_unresolved_relations()
585592
affected_entities: AffectedEntityIds = set()
@@ -593,12 +600,58 @@ async def resolve_project_relations(
593600
if not affected:
594601
break
595602

603+
remaining = await runtime.count_unresolved_relations()
604+
failed_entity_ids: set[EntityId] = set()
605+
deferred_entity_ids: set[EntityId] = set()
606+
607+
# Trigger: the resolution pass rewrote relation search rows for these sources
608+
# after their owners had already embedded them.
609+
# Why: chunk keys follow the row text, so those changed relation rows have no
610+
# current vector until something re-syncs them; nothing else revisits them.
611+
# Outcome: re-embed the affected sources after resolution committed. The batch
612+
# backend reports failed and deferred entities instead of raising, so drain
613+
# every deferred shard here; failed entities stay marked by the backend's
614+
# retry/readiness state and both sets are exposed on the result.
615+
if vector_sync is not None and affected_entities:
616+
pending_entity_ids = sorted(affected_entities)
617+
previous_pending_chunks: int | None = None
618+
while pending_entity_ids:
619+
sync_result = await vector_sync.sync_entity_vectors_batch(pending_entity_ids)
620+
failed_entity_ids.update(sync_result.failed_entity_ids)
621+
pending_entity_ids = list(sync_result.deferred_entity_ids)
622+
if not pending_entity_ids:
623+
break
624+
625+
pending_chunks = sync_result.chunks_total - sync_result.chunks_skipped
626+
if previous_pending_chunks is not None and pending_chunks >= previous_pending_chunks:
627+
logger.warning(
628+
"Stopped draining deferred relation vectors: no progress",
629+
deferred_entities=len(pending_entity_ids),
630+
pending_chunks=pending_chunks,
631+
)
632+
break
633+
previous_pending_chunks = pending_chunks
634+
635+
deferred_entity_ids.update(pending_entity_ids)
636+
if failed_entity_ids:
637+
logger.warning(
638+
"Relation vector sync left failed entities",
639+
failed_entities=len(failed_entity_ids),
640+
)
641+
if deferred_entity_ids:
642+
logger.warning(
643+
"Relation vector sync left deferred entities after drain",
644+
deferred_entities=len(deferred_entity_ids),
645+
)
646+
596647
result = ResolveRelationsResult(
597648
unresolved_before=unresolved_before,
598-
remaining=await runtime.count_unresolved_relations(),
649+
remaining=remaining,
599650
passes=passes,
600651
affected_entities=len(affected_entities),
601652
affected_entity_ids=frozenset(affected_entities),
653+
vector_sync_failed_entity_ids=frozenset(failed_entity_ids),
654+
vector_sync_deferred_entity_ids=frozenset(deferred_entity_ids),
602655
)
603656
logger.info(
604657
"Resolved project relations",
@@ -608,12 +661,4 @@ async def resolve_project_relations(
608661
passes=result.passes,
609662
affected_entities=result.affected_entities,
610663
)
611-
# Trigger: the resolution pass rewrote relation search rows for these sources
612-
# after their owners had already embedded them.
613-
# Why: chunk keys follow the row text, so those changed relation rows have no
614-
# current vector until something re-syncs them; nothing else revisits them.
615-
# Outcome: one batch re-embeds the affected sources after resolution committed;
616-
# a later batch or reindex is not required.
617-
if vector_sync is not None and affected_entities:
618-
await vector_sync.sync_entity_vectors_batch(sorted(affected_entities))
619664
return result

‎tests/index/test_local_inline_result_recorder.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from basic_memory.indexing.project_index_maintenance import ProjectIndexMovedEntitySearchRefresher
1515
from basic_memory.indexing.relation_resolution import RelationResolutionRuntime
1616
from basic_memory.runtime.projects import ProjectRuntimeReference
17+
from basic_memory.runtime.vector_sync import VectorSyncBatchResult
1718
from basic_memory.runtime.storage import (
1819
RuntimeStorageEventOperation,
1920
RuntimeStorageEventOperationKind,
@@ -58,7 +59,11 @@ def __init__(self) -> None:
5859

5960
async def sync_entity_vectors_batch(self, entity_ids, progress_callback=None):
6061
self.batches.append(list(entity_ids))
61-
return None
62+
return VectorSyncBatchResult(
63+
entities_total=len(entity_ids),
64+
entities_synced=len(entity_ids),
65+
entities_failed=0,
66+
)
6267

6368

6469
class StubRelationRuntime:

‎tests/indexing/test_relation_resolution.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,18 @@ async def sync_entity_vectors_batch(self, entity_ids: list[int]) -> VectorSyncBa
6868
)
6969

7070

71+
class ScriptedVectorSync:
72+
"""Returns scripted batch results and records requested entity batches."""
73+
74+
def __init__(self, results: list[VectorSyncBatchResult]) -> None:
75+
self.results = list(results)
76+
self.batches: list[list[int]] = []
77+
78+
async def sync_entity_vectors_batch(self, entity_ids: list[int]) -> VectorSyncBatchResult:
79+
self.batches.append(list(entity_ids))
80+
return self.results.pop(0)
81+
82+
7183
class FakeSession:
7284
created_count = 0
7385

@@ -541,6 +553,77 @@ async def test_resolve_project_relations_syncs_vectors_once_after_last_pass() ->
541553
assert vector_sync.batches == [[10, 11, 12]]
542554

543555

556+
@pytest.mark.asyncio
557+
async def test_resolve_project_relations_drains_deferred_vector_shards() -> None:
558+
runtime = StubRelationResolutionRuntime([3, 1], [{10, 11}, set()])
559+
vector_sync = ScriptedVectorSync(
560+
[
561+
VectorSyncBatchResult(
562+
entities_total=2,
563+
entities_synced=1,
564+
entities_failed=0,
565+
entities_deferred=1,
566+
deferred_entity_ids=(11,),
567+
chunks_total=2,
568+
chunks_skipped=1,
569+
),
570+
VectorSyncBatchResult(
571+
entities_total=1,
572+
entities_synced=1,
573+
entities_failed=0,
574+
chunks_total=1,
575+
chunks_skipped=1,
576+
),
577+
]
578+
)
579+
580+
result = await resolve_project_relations(runtime, vector_sync=vector_sync)
581+
582+
assert vector_sync.batches == [[10, 11], [11]]
583+
assert result.vector_sync_deferred_entity_ids == frozenset()
584+
585+
586+
@pytest.mark.asyncio
587+
async def test_resolve_project_relations_exposes_failed_vector_entities() -> None:
588+
runtime = StubRelationResolutionRuntime([3, 1], [{10, 11}, set()])
589+
vector_sync = ScriptedVectorSync(
590+
[
591+
VectorSyncBatchResult(
592+
entities_total=2,
593+
entities_synced=1,
594+
entities_failed=1,
595+
failed_entity_ids=(11,),
596+
)
597+
]
598+
)
599+
600+
result = await resolve_project_relations(runtime, vector_sync=vector_sync)
601+
602+
assert vector_sync.batches == [[10, 11]]
603+
assert result.vector_sync_failed_entity_ids == frozenset({11})
604+
assert result.vector_sync_deferred_entity_ids == frozenset()
605+
606+
607+
@pytest.mark.asyncio
608+
async def test_resolve_project_relations_stops_draining_without_progress() -> None:
609+
runtime = StubRelationResolutionRuntime([3, 1], [{10, 11}, set()])
610+
deferred = VectorSyncBatchResult(
611+
entities_total=1,
612+
entities_synced=0,
613+
entities_failed=0,
614+
entities_deferred=1,
615+
deferred_entity_ids=(11,),
616+
chunks_total=2,
617+
chunks_skipped=1,
618+
)
619+
vector_sync = ScriptedVectorSync([deferred, deferred])
620+
621+
result = await resolve_project_relations(runtime, vector_sync=vector_sync)
622+
623+
assert vector_sync.batches == [[10, 11], [11]]
624+
assert result.vector_sync_deferred_entity_ids == frozenset({11})
625+
626+
544627
@pytest.mark.asyncio
545628
async def test_resolve_project_relations_skips_vector_sync_when_nothing_resolved() -> None:
546629
runtime = StubRelationResolutionRuntime([1, 1], [set()])

0 commit comments

Comments
 (0)