feat(consolidation): parallelize llm_batch processing within a single op

[connorblack]

Summary

Replace the serial for llm_batch in llm_batches loop in run_consolidation_job with bounded asyncio.gather using a Semaphore at config.consolidation_llm_max_concurrent. The env knob has existed for a while (default 8) but was unused for single-bank workloads — the previous serial loop bottlenecked on single-LLM-call latency regardless of available capacity.

Why this matters

In production we observed consolidation_llm_max_concurrent=8 configured but the consolidator only ever made one in-flight LLM call per op. With Qwen3.6-35B-A3B-FP8 on a single GB10 GPU, latency per consolidation call is 5-10s once chat_template_kwargs.enable_thinking=false is set; with parallelism unlocked, the same wave processes 8 memories in roughly the same wall time as 1.

Verified on a 14k-memory bank: same-millisecond timestamp on the per-batch log lines for all batches in a wave (proves gather returned them as a wave) with no double-create or skipped failures across hundreds of batches.

Correctness invariants preserved

• Tag-group security boundary: memories with different tags still never share an LLM call (slicing happens upstream of the gather; each task operates on a single tag group's slice).
• Adaptive split-on-failure protocol: the inner pending queue (halve sub-batch on LLM failure, mark failed only at size=1) remains serial within each batch — that ordering is required by the split protocol.
• Disjoint memory sets per batch: LIMIT \$2 query produces a fixed set per fetch round; tag-group slicing produces non-overlapping subsets. No two batches in a gather wave can race on the same memory_id.
• DB commit semantics: cross-batch IDs aggregated into a single executemany per fetch round (functionally equivalent to the previous per-batch commits — IDs are disjoint).

Module surface added

• ConsolidationPerfLog.__iadd__ — accumulator (matches TokenUsage.__add__ codebase idiom).
• _BatchExecutionResult dataclass — self-contained per-batch outputs for parent aggregation.
• _record_result(stats, result) -> _ResultDelta — centralizes the action-vocabulary mapping (created/updated/merged/multiple/skipped/failed). Side benefit: pre-patch silently dropped merged actions from per-batch log lines; the new helper tracks them.
• _merge_pass_result(existing, new) -> dict — extracted from the dense inline block in the multi-pass loop; testable in isolation.
• _resolve_obs_tags_list(memory_tags, scope_spec) — translates observation_scopes spec into concrete tag-set passes; called once per llm_batch instead of once per sub_batch.
• _is_op_cancelled(memory_engine, operation_id) -> bool — predicate used at both the per-sub-batch check inside _execute_one_llm_batch and the end-of-wave checkpoint in run_consolidation_job.
• _execute_one_llm_batch(...) async helper extracted from the for-loop body. Accepts operation_id so per-sub-batch cancellation polling restores per-batch cancellation granularity (the gather wave checkpoint alone would have lengthened cancellation latency from ~5s to ~45s).

Per-batch logging

Each batch executes against its own ConsolidationPerfLog instance so timings/llm_calls/prompt_chars are not raced across concurrent gather participants. Per-task perf is merged into the parent via += after gather completes. The pre-patch snap-delta-style logging (capture-before, log-after) was incorrect under parallel execution; the new shape carries per-batch perf in the result dataclass and emits log lines deterministically ordered by batch_num.

Throughput notes

Empirical: ~16× speedup of consolidation throughput when combined with thinking-off (~5s/call vs ~80s/call). At a fixed concurrency of 8, the 8 in-flight LLM calls per wave saturate the GPU; cranking to 16 increases per-call latency due to KV-cache contention without proportional throughput gain (this is GPU-bound, not Python-bound). Conservative production setting: leave consolidation_llm_max_concurrent at the default of 8.

Test plan

• Local pytest suite at hindsight-api-slim/tests/test_consolidation*.py requires HINDSIGHT_API_LLM_API_KEY for testcontainers setup; was not exercised against this fork. CI's core lane (paths-filter: hindsight-api-slim/**) will run the suite end-to-end.
• Production smoke verified on a 14k-memory bank: zero 404s, zero schema rejections, zero duplicate observations, zero consolidation_failed_at regressions across 1000+ batches.
• pre-commit hooks (ruff check --fix, ruff format, ty check) all green per scripts/hooks/lint.sh.

Notes for reviewers

Three commits on the branch — visible iteration via two rounds of internal review (reuse / quality / efficiency). Squash on merge is fine if that's the project preference; the final state is what matters.

Cancellation latency under parallelism: per-sub-batch via operation_id (~one DB roundtrip per sub_batch) plus end-of-wave checkpoint. With wave size ≤ consolidation_batch_size (default 50) and waves completing in seconds with thinking off, worst-case cancellation latency is bounded by the slowest LLM call in the active wave.

Related follow-ups (not in this PR)

• Inter-wave fetch pipelining: the next DB fetch round can't start until the current gather wave completes. With <0.1% inter-wave overhead in practice (5ms fetch vs ~6s wave) this is negligible at default config; would matter only on much smaller round sizes. Worth filing as a separate issue if motivated by observed throughput pain.

[Copilot]

Pull request overview

This PR parallelizes consolidation LLM-batch execution within each DB fetch “wave” by replacing a serial per-batch loop with bounded asyncio.gather() concurrency, while refactoring the batch execution into an async helper and making per-batch performance/log aggregation safe under parallelism.

Changes:

• Execute llm_batches concurrently with a per-wave asyncio.Semaphore bounded by config.consolidation_llm_max_concurrent.
• Extract per-batch execution into _execute_one_llm_batch() and return a _BatchExecutionResult for parent aggregation.
• Centralize stats/log accounting via _record_result() / _ResultDelta and merge per-task perf into the parent via ConsolidationPerfLog.__iadd__.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[nicoloboschi]

Must fix: same-tag batches must NOT run in parallel

The gather parallelizes all llm_batches, but multiple batches can come from the same tag group (when a group has more memories than llm_batch_size):

for group in tag_groups.values():
    for i in range(0, len(group), llm_batch_size):
        llm_batches.append(group[i : i + llm_batch_size])  # multiple batches, same tags

When two same-tag batches run concurrently, both read the observation network at the same time — so batch 2 never sees observations created/updated by batch 1. The LLM decides "create new observation" in both batches for what should have been a single consolidated observation. Memories that should be consolidated together won't be, because each batch operates on a stale snapshot of the observation graph.

Suggested fix

Parallelize across tag groups (disjoint observation scopes → safe), but keep batches within a tag group serial (shared observation scope → must see prior writes):

async def _process_tag_group(group_batches, sem, ...):
    results = []
    async with sem:  # bound concurrency across all tag groups
        for batch in group_batches:
            results.append(await _execute_one_llm_batch(batch, ...))
    return results

# Build per-tag-group batch lists instead of flattening
per_tag_group: list[list[...]] = []
for group in tag_groups.values():
    per_tag_group.append([group[i:i+llm_batch_size] for i in range(0, len(group), llm_batch_size)])

gather_results = await asyncio.gather(
    *(_process_tag_group(batches, sem, ...) for batches in per_tag_group),
    return_exceptions=True,
)

Test gap

test_consolidation_honors_max_concurrent uses distinct tags per memory (person0, person1, ...) so it can't catch this. A regression test should ingest N memories with the same tag, set llm_batch_size small enough to split them, and assert that batch 2's observation recall sees batch 1's writes (i.e., observations are consolidated together, not duplicated).