diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 08876f7b92..b47a9c47f9 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -273,6 +273,12 @@ def _split_contents_into_sub_batches( would pass through as a ``1/1`` sub-batch holding the entire payload — which contradicts the splitter's log and lets the orchestrator OOM under realistic memory limits (see issue #1571). + + Used by the in-process ``retain_batch_async`` path, which processes + the returned sub-batches SEQUENTIALLY with ``is_first_batch=(i==1)``. + The async submission path uses ``_split_contents_into_async_children`` + instead, which never fragments a single item across children — see + that helper for the reasoning. """ from .retain import fact_extraction @@ -326,6 +332,66 @@ def _flush() -> None: return _SubBatchSplit(sub_batches=sub_batches, origin_indices=origin_indices) +def _split_contents_into_async_children( + contents: list[RetainContentDict], + tokens_per_batch: int, +) -> list[list[RetainContentDict]]: + """Pack retain contents into child operations for async submission. + + Unlike ``_split_contents_into_sub_batches`` (used by the in-process + path), this NEVER fragments a single input item across multiple + children. Items where ``count_tokens(content) > tokens_per_batch`` + are emitted as their own single-item child holding the FULL + un-chunked content; the in-process ``retain_batch_async`` then + re-chunks them SEQUENTIALLY inside one worker slot with correct + ``is_first_batch=(i==1)`` semantics. + + The previous behavior — chunking oversized items into N independent + child async-operations sharing one ``document_id`` — let workers + claim siblings concurrently with no per-document gate (the busy-bank + guard in ``claim_tasks`` only covers consolidation). Each concurrent + child ran ``handle_document_tracking(is_first_batch=True)``, which + cascade-deletes the prior winner's ``memory_units`` for that + document. The loser's final ANN pass then attempted to insert + ``memory_links`` referencing now-deleted units → FK violations on + ``fk_memory_links_from_unit_id_memory_units``, partial document + state, and worker thread exhaustion from sentence-transformer pools + spun up per concurrent child. See issue #1795. + + Items smaller than the budget are still packed together so genuinely + independent items keep cross-worker parallelism. + """ + children: list[list[RetainContentDict]] = [] + current: list[RetainContentDict] = [] + current_tokens = 0 + + def _flush() -> None: + nonlocal current, current_tokens + if current: + children.append(current) + current = [] + current_tokens = 0 + + for item in contents: + item_tokens = count_tokens(item.get("content", "") or "") + + if item_tokens > tokens_per_batch: + # Oversized: flush in-flight items into their own child, + # then emit this item AS-IS (un-chunked) into its own child. + # The worker will sequentially chunk it inside retain_batch_async. + _flush() + children.append([item]) + continue + + if current and current_tokens + item_tokens > tokens_per_batch: + _flush() + current.append(item) + current_tokens += item_tokens + + _flush() + return children + + def _is_invalid_embedding_dimension_error(e: Exception) -> bool: """Return True for deterministic embedding-dimension failures. @@ -9655,14 +9721,14 @@ async def submit_async_retain( config = get_config() tokens_per_batch = config.retain_batch_tokens - # Split into sub-batches based on token count. Oversized single - # items get chunked into per-chunk sub-batches by the helper - # (see issue #1571). origin_indices is unused here because - # submit_async_retain returns only the parent operation_id, not - # per-input results. - sub_batches = _split_contents_into_sub_batches( - cast(list[RetainContentDict], contents), tokens_per_batch - ).sub_batches + # Pack items into child operations by token budget. An oversized + # single item is emitted as its own un-chunked child rather than + # being fragmented across siblings — workers have no + # per-document serialization, so concurrent siblings would race + # on the same document_id and trigger FK violations in the final + # ANN pass (issue #1795). The worker's in-process splitter + # handles intra-document chunking sequentially. + sub_batches = _split_contents_into_async_children(cast(list[RetainContentDict], contents), tokens_per_batch) # Log splitting info if we actually split if len(sub_batches) > 1: @@ -9670,13 +9736,13 @@ async def submit_async_retain( if len(sub_batches) <= 20: logger.info( f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). " - f"Split into {len(sub_batches)} sub-batches: {sub_batch_sizes} items each" + f"Split into {len(sub_batches)} child operations: {sub_batch_sizes} items each" ) else: logger.info( f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). " - f"Split into {len(sub_batches)} sub-batches " - f"(items per sub-batch: min={min(sub_batch_sizes)}, " + f"Split into {len(sub_batches)} child operations " + f"(items per child: min={min(sub_batch_sizes)}, " f"max={max(sub_batch_sizes)}, total={sum(sub_batch_sizes)})" ) @@ -9742,7 +9808,7 @@ async def submit_async_retain( if len(sub_batches) > 1: sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch) logger.info( - f"Submitting sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens" + f"Submitting child {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens" ) task_payload: dict[str, Any] = {"contents": sub_batch} diff --git a/hindsight-api-slim/tests/test_async_batch_retain.py b/hindsight-api-slim/tests/test_async_batch_retain.py index 2771e992d0..78f421bd28 100644 --- a/hindsight-api-slim/tests/test_async_batch_retain.py +++ b/hindsight-api-slim/tests/test_async_batch_retain.py @@ -832,6 +832,173 @@ async def failing_submit_task(_task_dict): assert payload["contents"] == [{"content": "hello", "document_id": "d1"}] +# --------------------------------------------------------------------------- +# Regression tests for issue #1795: submit_async_retain must NOT fragment a +# single oversized item across multiple child async-operations. Workers have +# no per-document serialization for retain (the busy-bank guard in +# claim_tasks only covers consolidation), so concurrent siblings sharing one +# document_id race on handle_document_tracking(is_first_batch=True), cascade- +# delete each other's memory_units, and trip FK violations on memory_links in +# the final ANN pass. The fix: keep the oversized item un-chunked in a single +# child; the worker's in-process splitter handles intra-document chunking +# sequentially with correct is_first_batch semantics. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_oversized_single_item_creates_one_child_not_many(memory, request_context, monkeypatch): + """One oversized document → one child operation holding the un-chunked content.""" + from hindsight_api.config import get_config + from hindsight_api.engine.memory_engine import count_tokens + + bank_id = f"test_oversized_one_child_{uuid.uuid4().hex[:8]}" + pool = await memory._get_pool() + await _ensure_bank(pool, bank_id) + + # No-op the worker dispatch: this is a structural assertion on the + # async_operations rows that submit_async_retain inserts (those rows are + # committed before submit_task is invoked), so we don't need the + # SyncTaskBackend to drive the slow LLM pipeline to completion. + async def noop_submit_task(_task_dict): + return None + + monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task) + + tokens_per_batch = get_config().retain_batch_tokens + # Build content that comfortably exceeds the per-batch token budget. + # The pre-fix code would split this into many siblings, all sharing + # the same document_id and racing on handle_document_tracking. + unit = "The quick brown fox jumps over the lazy dog. " * 50 # ~570 tokens + repetitions = max(1, (tokens_per_batch * 4) // count_tokens(unit) + 1) + big_content = unit * repetitions + document_id = f"doc-oversize-{uuid.uuid4().hex[:8]}" + + total_tokens = count_tokens(big_content) + assert total_tokens > tokens_per_batch, ( + f"Test setup error: content has {total_tokens} tokens, must exceed the batch budget of {tokens_per_batch}" + ) + + result = await memory.submit_async_retain( + bank_id=bank_id, + contents=[{"content": big_content, "document_id": document_id}], + request_context=request_context, + ) + parent_operation_id = result["operation_id"] + + # Parent metadata must say one child — even though token count is huge, + # the un-fragmenting splitter keeps the single item in one child. + parent_row = await pool.fetchrow( + "SELECT result_metadata FROM async_operations WHERE operation_id = $1", + uuid.UUID(parent_operation_id), + ) + parent_meta = ( + json.loads(parent_row["result_metadata"]) + if isinstance(parent_row["result_metadata"], str) + else parent_row["result_metadata"] + ) + assert parent_meta["num_sub_batches"] == 1, ( + f"Expected 1 child for an oversized single item, got {parent_meta['num_sub_batches']}. " + "Issue #1795: per-chunk children race on the shared document_id." + ) + + # Exactly one retain child row. + children = await pool.fetch( + """ + SELECT operation_id, status, task_payload, result_metadata + FROM async_operations + WHERE bank_id = $1 AND operation_type = 'retain' + """, + bank_id, + ) + assert len(children) == 1, ( + f"Expected exactly 1 child retain operation for bank_id={bank_id}, " + f"got {len(children)}. Issue #1795 regression: oversized item was " + f"fragmented into per-chunk children." + ) + + child = children[0] + payload = child["task_payload"] + payload = json.loads(payload) if isinstance(payload, str) else payload + assert payload["type"] == "batch_retain" + assert payload["bank_id"] == bank_id + # Child holds the full un-chunked content — the worker's in-process + # splitter will re-chunk it sequentially with is_first_batch=(i==1). + assert len(payload["contents"]) == 1, ( + f"Child payload should hold the original single item, got {len(payload['contents'])} items" + ) + assert payload["contents"][0]["document_id"] == document_id + assert payload["contents"][0]["content"] == big_content, ( + "Child must carry the FULL un-chunked content — chunking happens inside the worker, not at submit time" + ) + + +@pytest.mark.asyncio +async def test_oversized_item_among_small_items_keeps_small_items_packed(memory, request_context, monkeypatch): + """Mixed batch: small items pack together; oversized item isolates to its own child.""" + from hindsight_api.config import get_config + from hindsight_api.engine.memory_engine import count_tokens + + bank_id = f"test_mixed_pack_{uuid.uuid4().hex[:8]}" + pool = await memory._get_pool() + await _ensure_bank(pool, bank_id) + + # Structural assertion on async_operations rows only — skip worker dispatch. + async def noop_submit_task(_task_dict): + return None + + monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task) + + tokens_per_batch = get_config().retain_batch_tokens + big_unit = "The quick brown fox jumps over the lazy dog. " * 50 + big_repetitions = max(1, (tokens_per_batch * 3) // count_tokens(big_unit) + 1) + big_content = big_unit * big_repetitions + + big_doc = f"doc-big-{uuid.uuid4().hex[:8]}" + small_docs = [f"doc-small-{i}-{uuid.uuid4().hex[:6]}" for i in range(3)] + contents = [ + {"content": "Alice works at Google.", "document_id": small_docs[0]}, + {"content": big_content, "document_id": big_doc}, + {"content": "Bob loves Python.", "document_id": small_docs[1]}, + {"content": "Carol writes Rust.", "document_id": small_docs[2]}, + ] + + result = await memory.submit_async_retain( + bank_id=bank_id, + contents=contents, + request_context=request_context, + ) + + children = await pool.fetch( + """ + SELECT task_payload + FROM async_operations + WHERE bank_id = $1 AND operation_type = 'retain' + ORDER BY created_at, operation_id + """, + bank_id, + ) + + # Find the child holding the big doc — it must hold ONLY the big doc. + big_child = None + small_doc_ids_seen: list[str] = [] + for row in children: + payload = row["task_payload"] + payload = json.loads(payload) if isinstance(payload, str) else payload + items = payload["contents"] + doc_ids = [item["document_id"] for item in items] + if big_doc in doc_ids: + assert big_child is None, f"Big doc {big_doc} should appear in exactly one child, found in 2" + big_child = items + assert doc_ids == [big_doc], f"Child holding the oversized item must hold ONLY that item, got {doc_ids}" + assert items[0]["content"] == big_content + else: + small_doc_ids_seen.extend(doc_ids) + + assert big_child is not None, "Big doc must appear in some child" + # Every small input present, exactly once, across the other children. + assert sorted(small_doc_ids_seen) == sorted(small_docs) + + @pytest.mark.asyncio async def test_submit_async_batch_retain_rolls_back_parent_on_child_failure( memory_no_llm_verify, request_context, monkeypatch diff --git a/hindsight-api-slim/tests/test_batch_chunking.py b/hindsight-api-slim/tests/test_batch_chunking.py index 60674b5c75..f0f94d2827 100644 --- a/hindsight-api-slim/tests/test_batch_chunking.py +++ b/hindsight-api-slim/tests/test_batch_chunking.py @@ -7,6 +7,7 @@ from hindsight_api import MemoryEngine from hindsight_api.engine.memory_engine import ( + _split_contents_into_async_children, _split_contents_into_sub_batches, count_tokens, ) @@ -118,6 +119,158 @@ def test_split_small_batch_returns_single_sub_batch(): assert split.origin_indices == [[0, 1]] +# --------------------------------------------------------------------------- +# Regression tests for issue #1795: the async-submit splitter must NEVER +# fragment a single oversized item across multiple children. Each child +# becomes an independent async_operations row claimed by workers in +# parallel with no per-document gate, so siblings would race on the same +# document_id, cascade-delete each other's memory_units, and trip FK +# violations on memory_links in the final ANN pass. +# --------------------------------------------------------------------------- + + +def test_async_children_oversized_single_item_stays_in_one_child(): + """An oversized single item must NOT be split across children.""" + tokens_per_batch = 1_000 + big_content = "The quick brown fox jumps over the lazy dog. " * 1_000 + item = {"content": big_content, "document_id": "doc-oversize"} + assert count_tokens(big_content) > tokens_per_batch + + children = _split_contents_into_async_children([item], tokens_per_batch) + + # Exactly one child holding the full un-chunked item — the worker + # will sequentially chunk it inside retain_batch_async. + assert len(children) == 1, ( + f"Oversized single item must become exactly one child, got {len(children)}. " + "Regressing to per-chunk children causes the issue #1795 race." + ) + assert children[0] == [item] + assert children[0][0]["content"] == big_content + + +def test_async_children_oversized_item_preserves_metadata(): + """The single-child payload must carry every original key untouched.""" + tokens_per_batch = 500 + item = { + "content": "Alice met Bob at the coffee shop. " * 500, + "document_id": "doc-42", + "context": "shared-context", + "tags": ["t1", "t2"], + "metadata": {"source": "test"}, + } + + children = _split_contents_into_async_children([item], tokens_per_batch) + + assert len(children) == 1 + assert children[0] == [item] + + +def test_async_children_packs_small_items_by_budget(): + """Small items pack into shared children up to the token budget.""" + tokens_per_batch = 100 + item_text = "Alice works at Google and Bob loves Python programming language." + item_tokens = count_tokens(item_text) + # Pick enough items to force at least 2 children at this budget. + num_items = max(4, (tokens_per_batch // max(item_tokens, 1)) * 3) + contents = [{"content": item_text, "document_id": f"doc-{i}"} for i in range(num_items)] + total = sum(count_tokens(c["content"]) for c in contents) + assert total > tokens_per_batch, ( + f"Test setup error: {total} tokens does not exceed budget {tokens_per_batch}" + ) + + children = _split_contents_into_async_children(contents, tokens_per_batch) + + assert len(children) >= 2 + # Every child stays within budget. + for child in children: + child_tokens = sum(count_tokens(item["content"]) for item in child) + assert child_tokens <= tokens_per_batch, ( + f"Child holding {child_tokens} tokens exceeds budget {tokens_per_batch}" + ) + # Every input item appears exactly once across all children. + flat = [item for child in children for item in child] + assert flat == contents + + +def test_async_children_mixed_small_and_oversized(): + """Small items pack together; the oversized item gets its own child as-is.""" + tokens_per_batch = 1_000 + small_a = {"content": "Alice works at Google. " * 5, "document_id": "doc-a"} + big = {"content": "The quick brown fox jumps over the lazy dog. " * 1_000, "document_id": "doc-big"} + small_b = {"content": "Bob loves Python. " * 5, "document_id": "doc-b"} + small_c = {"content": "Carol writes Rust. " * 5, "document_id": "doc-c"} + + children = _split_contents_into_async_children([small_a, big, small_b, small_c], tokens_per_batch) + + # The big item must be in its own child, un-chunked. + big_children = [c for c in children if any(item is big or item.get("document_id") == "doc-big" for item in c)] + assert len(big_children) == 1, ( + f"Oversized item must occupy exactly one child, found {len(big_children)} containing it" + ) + assert big_children[0] == [big] + assert big_children[0][0]["content"] == big["content"] + + # Every small item is present somewhere. + flat_doc_ids = [item["document_id"] for child in children for item in child] + for small in (small_a, small_b, small_c): + assert small["document_id"] in flat_doc_ids + + # No child fragments the big item. + for child in children: + big_count = sum(1 for item in child if item.get("document_id") == "doc-big") + assert big_count <= 1 + + +def test_async_children_each_oversized_item_gets_own_child(): + """Multiple oversized items → one child per item, each un-chunked.""" + tokens_per_batch = 500 + big_a = {"content": "Alpha sentence. " * 500, "document_id": "doc-a"} + big_b = {"content": "Bravo sentence. " * 500, "document_id": "doc-b"} + big_c = {"content": "Charlie sentence. " * 500, "document_id": "doc-c"} + + children = _split_contents_into_async_children([big_a, big_b, big_c], tokens_per_batch) + + assert len(children) == 3 + assert children == [[big_a], [big_b], [big_c]] + + +def test_async_children_single_small_item_returns_one_child(): + """A single small item is one child of one item — no splitting.""" + tokens_per_batch = 10_000 + item = {"content": "Alice works at Google", "document_id": "doc-1"} + + children = _split_contents_into_async_children([item], tokens_per_batch) + + assert children == [[item]] + + +def test_async_children_empty_returns_empty(): + """Empty input → empty output (no spurious child).""" + assert _split_contents_into_async_children([], 1_000) == [] + + +def test_async_children_oversized_at_boundaries(): + """Oversized items at the start, middle, and end all isolate correctly.""" + tokens_per_batch = 1_000 + small_a = {"content": "Small A. " * 5, "document_id": "doc-sa"} + small_b = {"content": "Small B. " * 5, "document_id": "doc-sb"} + big_start = {"content": "First big. " * 1_000, "document_id": "doc-bs"} + big_end = {"content": "Last big. " * 1_000, "document_id": "doc-be"} + + children = _split_contents_into_async_children([big_start, small_a, small_b, big_end], tokens_per_batch) + + # First child holds big_start alone; last child holds big_end alone. + assert children[0] == [big_start] + assert children[-1] == [big_end] + # Middle children contain only small items. + middle_items = [item for child in children[1:-1] for item in child] + assert all(item["document_id"] in {"doc-sa", "doc-sb"} for item in middle_items) + # Every small input is present. + flat_ids = [item["document_id"] for child in children for item in child] + assert "doc-sa" in flat_ids + assert "doc-sb" in flat_ids + + @pytest.mark.asyncio async def test_large_batch_auto_chunks(memory, request_context): bank_id = "test_chunking_agent"