From ba172b821d5b985315b366a82ef94e344f19080e Mon Sep 17 00:00:00 2001 From: Connor Black Date: Tue, 12 May 2026 19:49:07 -0400 Subject: [PATCH 1/4] feat(consolidation): parallelize llm_batch processing within a single op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces 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. Each batch executes in its own ConsolidationPerfLog instance so timings/llm_calls/prompt_chars are not raced across concurrent gather participants; the per-task perf is merged into the parent's shared perf after gather completes. Why this matters: - consolidation_llm_max_concurrent (default 8) was unused for single-bank workloads — the previous serial loop bottlenecked on single-LLM-call latency. With gather + Semaphore we saturate up to N parallel LLM calls within one op for an N x speedup. - The tag-group security boundary (memories with different tags never share an LLM call) is preserved unchanged. - The adaptive split-on-failure protocol remains serial within each batch (correctness requirement of the split protocol). - The cross-batch DB commit is now a single round-trip per fetch wave instead of one commit per batch (functionally equivalent — IDs are disjoint by tag-grouping). Cancellation granularity changes from per-batch to per-gather-wave; acceptable since waves are bounded by consolidation_batch_size memories (default 50) and complete in seconds. New module surface: - ConsolidationPerfLog.merge() — aggregate per-task perf into shared - _BatchExecutionResult dataclass — self-contained per-batch result - _execute_one_llm_batch() — extracted per-batch work, used as gather task --- .../engine/consolidation/consolidator.py | 466 +++++++++++------- 1 file changed, 294 insertions(+), 172 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index cd81e5ab88..261bd1e68c 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -15,6 +15,7 @@ - Mental models: user-defined queries stored in the mental_models table, refreshed on demand via reflect """ +import asyncio import json import logging import time @@ -215,6 +216,21 @@ def record_llm_call(self, obs_count: int, prompt_chars: int) -> None: self.total_obs_in_context += obs_count self.total_prompt_chars += prompt_chars + def merge(self, other: "ConsolidationPerfLog") -> None: + """Merge another perf log's contributions into this one. + + Used to safely aggregate per-batch perf state collected by concurrent + consolidation tasks (see _execute_one_llm_batch). Each task gets its own + ConsolidationPerfLog so timings/llm_calls/prompt_chars are not racing + on the shared parent perf during asyncio.gather, then merged in-order + after gather completes. + """ + for key, val in other.timings.items(): + self.timings[key] = self.timings.get(key, 0) + val + self.llm_calls += other.llm_calls + self.total_obs_in_context += other.total_obs_in_context + self.total_prompt_chars += other.total_prompt_chars + def flush(self) -> None: """Flush all log lines to the logger.""" total_time = time.time() - self.start_time @@ -225,6 +241,193 @@ def flush(self) -> None: logger.info(log_output) +@dataclass +class _BatchExecutionResult: + """Self-contained result of one llm_batch execution. + + Returned by ``_execute_one_llm_batch`` so the parent ``run_consolidation_job`` + can aggregate stats, perf, and DB writes from concurrent batch tasks after + asyncio.gather completes. Each field is independent — no shared mutable state + is touched during batch execution; aggregation is single-threaded after gather. + """ + + batch_num: int + memories_count: int + succeeded_ids: list[Any] + failed_ids: list[Any] + results: list[dict[str, Any]] + deleted: int + tags: set[str] + perf: ConsolidationPerfLog + elapsed: float + + +async def _execute_one_llm_batch( + *, + bank_id: str, + llm_batch: list[dict[str, Any]], + llm_batch_num: int, + pool: Any, + memory_engine: "MemoryEngine", + llm_config: Any, + request_context: "RequestContext", + config: Any, +) -> _BatchExecutionResult: + """Execute one llm_batch with adaptive split-on-failure, returning a self-contained result. + + Designed to be safely run concurrently with other batches via asyncio.gather + under a Semaphore at config.consolidation_llm_max_concurrent. All mutable + state (perf, succeeded_ids, failed_ids, tags) is local to this invocation; + nothing is shared with other concurrent invocations. The caller merges + these results into shared state after gather completes. + + The adaptive split logic (halve sub-batch on LLM failure, retry, mark failed + only at sub_batch=1) remains serial within this batch — that ordering is a + correctness requirement of the split protocol. + """ + batch_start = time.time() + batch_perf = ConsolidationPerfLog(bank_id) + + all_results: list[dict[str, Any]] = [] + all_deleted = 0 + succeeded_ids: list[Any] = [] + failed_ids: list[Any] = [] + batch_tags: set[str] = set() + + # Track tags for mental model refresh filtering (per-batch local set) + for memory in llm_batch: + memory_tags = memory.get("tags") or [] + if memory_tags: + batch_tags.update(memory_tags) + + # Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch + # and retry, down to batch_size=1. Only if a single-memory batch still fails is + # the memory marked with consolidation_failed_at and excluded from future runs + # until explicitly retried via the API. + pending: list[list[dict[str, Any]]] = [llm_batch] + while pending: + sub_batch = pending.pop(0) + + async with acquire_with_retry(pool) as conn: + # Determine observation_scopes for this sub-batch. All memories share + # the same tags (enforced by tag_groups), so we only check the first memory. + # asyncpg returns JSONB columns as raw JSON strings, so parse if needed. + _obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None + _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw + + # Resolve the scope spec into a concrete list[list[str]] (or None for combined). + if _obs_parsed == "per_tag": + _memory_tags = sub_batch[0].get("tags") or [] + obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None + elif _obs_parsed == "all_combinations": + _memory_tags = sub_batch[0].get("tags") or [] + obs_tags_list = ( + [ + list(combo) + for r in range(1, len(_memory_tags) + 1) + for combo in combinations(_memory_tags, r) + ] + if _memory_tags + else None + ) + elif _obs_parsed == "combined" or _obs_parsed is None: + obs_tags_list = None # single combined pass (default behaviour) + else: + # explicit list[list[str]] + obs_tags_list = _obs_parsed + + sub_deleted: int = 0 + sub_llm_failed = False + if obs_tags_list: + # Multi-pass: run one observation consolidation pass per tag set + sub_results: list[dict[str, Any]] = [] + for obs_tags in obs_tags_list: + pass_results, pass_deleted, pass_failed = await _process_memory_batch( + conn=conn, + memory_engine=memory_engine, + llm_config=llm_config, + bank_id=bank_id, + memories=sub_batch, + request_context=request_context, + perf=batch_perf, + config=config, + obs_tags_override=obs_tags, + ) + sub_deleted += pass_deleted + sub_llm_failed = sub_llm_failed or pass_failed + # Merge results: prefer non-skipped actions + if not sub_results: + sub_results = pass_results + else: + for i, (existing, new) in enumerate(zip(sub_results, pass_results)): + if existing.get("action") == "skipped" and new.get("action") != "skipped": + sub_results[i] = new + elif existing.get("action") != "skipped" and new.get("action") != "skipped": + # Both did something — combine into "multiple" + existing_created = existing.get( + "created", 1 if existing.get("action") == "created" else 0 + ) + existing_updated = existing.get( + "updated", 1 if existing.get("action") == "updated" else 0 + ) + new_created = new.get("created", 1 if new.get("action") == "created" else 0) + new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0) + total = existing_created + existing_updated + new_created + new_updated + sub_results[i] = { + "action": "multiple", + "created": existing_created + new_created, + "updated": existing_updated + new_updated, + "merged": 0, + "total_actions": total, + } + else: + # Normal single pass using the memory's own tags + sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch( + conn=conn, + memory_engine=memory_engine, + llm_config=llm_config, + bank_id=bank_id, + memories=sub_batch, + request_context=request_context, + perf=batch_perf, + config=config, + ) + + all_deleted += sub_deleted + + if sub_llm_failed and len(sub_batch) > 1: + # Split and retry with smaller batches + mid = len(sub_batch) // 2 + logger.warning( + f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)}," + f" splitting into {mid}/{len(sub_batch) - mid}" + ) + pending[0:0] = [sub_batch[:mid], sub_batch[mid:]] + elif sub_llm_failed: + # batch_size=1 and still failing — mark as permanently failed for now + failed_ids.append(sub_batch[0]["id"]) + all_results.append({"action": "failed"}) + logger.warning( + f"[CONSOLIDATION] bank={bank_id} LLM failed for single memory" + f" {sub_batch[0]['id']}, marking consolidation_failed_at" + ) + else: + succeeded_ids.extend(m["id"] for m in sub_batch) + all_results.extend(sub_results) + + return _BatchExecutionResult( + batch_num=llm_batch_num, + memories_count=len(llm_batch), + succeeded_ids=succeeded_ids, + failed_ids=failed_ids, + results=all_results, + deleted=all_deleted, + tags=batch_tags, + perf=batch_perf, + elapsed=time.time() - batch_start, + ) + + async def run_consolidation_job( memory_engine: "MemoryEngine", bank_id: str, @@ -363,166 +566,63 @@ async def run_consolidation_job( for i in range(0, len(group), llm_batch_size): llm_batches.append(group[i : i + llm_batch_size]) - for llm_batch in llm_batches: - llm_batch_num += 1 - llm_batch_start = time.time() - - # Snapshot perf and stats before this LLM batch - snap_timings = perf.timings.copy() - snap_llm_calls = perf.llm_calls - snap_total_chars = perf.total_prompt_chars - snap_stats = stats.copy() - - # Track tags for mental model refresh filtering - for memory in llm_batch: - memory_tags = memory.get("tags") or [] - if memory_tags: - consolidated_tags.update(memory_tags) - - # Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch - # and retry, down to batch_size=1. Only if a single-memory batch still fails is - # the memory marked with consolidation_failed_at and excluded from future runs - # until explicitly retried via the API. - all_results: list[dict[str, Any]] = [] - all_deleted = 0 - succeeded_ids: list[Any] = [] - failed_ids: list[Any] = [] - - pending: list[list[dict[str, Any]]] = [llm_batch] - while pending: - sub_batch = pending.pop(0) - - async with acquire_with_retry(pool) as conn: - # Determine observation_scopes for this sub-batch. All memories share - # the same tags (enforced by tag_groups), so we only check the first memory. - # asyncpg returns JSONB columns as raw JSON strings, so parse if needed. - _obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None - _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw - - # Resolve the scope spec into a concrete list[list[str]] (or None for combined). - if _obs_parsed == "per_tag": - _memory_tags = sub_batch[0].get("tags") or [] - obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None - elif _obs_parsed == "all_combinations": - _memory_tags = sub_batch[0].get("tags") or [] - obs_tags_list = ( - [ - list(combo) - for r in range(1, len(_memory_tags) + 1) - for combo in combinations(_memory_tags, r) - ] - if _memory_tags - else None - ) - elif _obs_parsed == "combined" or _obs_parsed is None: - obs_tags_list = None # single combined pass (default behaviour) - else: - # explicit list[list[str]] - obs_tags_list = _obs_parsed - - sub_deleted: int = 0 - sub_llm_failed = False - if obs_tags_list: - # Multi-pass: run one observation consolidation pass per tag set - sub_results: list[dict[str, Any]] = [] - for obs_tags in obs_tags_list: - pass_results, pass_deleted, pass_failed = await _process_memory_batch( - conn=conn, - memory_engine=memory_engine, - llm_config=llm_config, - bank_id=bank_id, - memories=sub_batch, - request_context=request_context, - perf=perf, - config=config, - obs_tags_override=obs_tags, - ) - sub_deleted += pass_deleted - sub_llm_failed = sub_llm_failed or pass_failed - # Merge results: prefer non-skipped actions - if not sub_results: - sub_results = pass_results - else: - for i, (existing, new) in enumerate(zip(sub_results, pass_results)): - if existing.get("action") == "skipped" and new.get("action") != "skipped": - sub_results[i] = new - elif existing.get("action") != "skipped" and new.get("action") != "skipped": - # Both did something — combine into "multiple" - existing_created = existing.get( - "created", 1 if existing.get("action") == "created" else 0 - ) - existing_updated = existing.get( - "updated", 1 if existing.get("action") == "updated" else 0 - ) - new_created = new.get("created", 1 if new.get("action") == "created" else 0) - new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0) - total = existing_created + existing_updated + new_created + new_updated - sub_results[i] = { - "action": "multiple", - "created": existing_created + new_created, - "updated": existing_updated + new_updated, - "merged": 0, - "total_actions": total, - } - else: - # Normal single pass using the memory's own tags - sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch( - conn=conn, - memory_engine=memory_engine, - llm_config=llm_config, - bank_id=bank_id, - memories=sub_batch, - request_context=request_context, - perf=perf, - config=config, - ) - - all_deleted += sub_deleted - - if sub_llm_failed and len(sub_batch) > 1: - # Split and retry with smaller batches - mid = len(sub_batch) // 2 - logger.warning( - f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)}," - f" splitting into {mid}/{len(sub_batch) - mid}" - ) - pending[0:0] = [sub_batch[:mid], sub_batch[mid:]] - elif sub_llm_failed: - # batch_size=1 and still failing — mark as permanently failed for now - failed_ids.append(sub_batch[0]["id"]) - all_results.append({"action": "failed"}) - logger.warning( - f"[CONSOLIDATION] bank={bank_id} LLM failed for single memory" - f" {sub_batch[0]['id']}, marking consolidation_failed_at" - ) - else: - succeeded_ids.extend(m["id"] for m in sub_batch) - all_results.extend(sub_results) - - # Commit consolidated_at / consolidation_failed_at in a single DB round-trip - async with acquire_with_retry(pool) as conn: - if succeeded_ids: - await conn.executemany( - f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1", - [(mem_id,) for mem_id in succeeded_ids], - ) - if failed_ids: - await conn.executemany( - f"UPDATE {fq_table('memory_units')} SET consolidation_failed_at = NOW() WHERE id = $1", - [(mem_id,) for mem_id in failed_ids], - ) + # Parallel batch execution with semaphore at consolidation_llm_max_concurrent. + # Each llm_batch operates on disjoint memories (the outer SELECT … LIMIT + # query returns a fixed set this round; tag-group slicing produces + # non-overlapping subsets), so concurrent execution is safe — no two + # batches can touch the same memory_id. The tag-group security boundary + # is preserved unchanged: cross-tag memories never share an LLM call. + # + # Pre-refactor (serial for-loop) left consolidation_llm_max_concurrent + # unused for single-bank workloads, bottlenecking on single-LLM-call + # latency. With asyncio.gather + Semaphore we saturate up to N parallel + # LLM calls within one consolidation op for an N× throughput gain. + sem = asyncio.Semaphore(max(1, config.consolidation_llm_max_concurrent)) + batch_base_num = llm_batch_num + + async def _bounded_batch(idx: int, b: list[dict[str, Any]]) -> _BatchExecutionResult: + async with sem: + return await _execute_one_llm_batch( + bank_id=bank_id, + llm_batch=b, + llm_batch_num=batch_base_num + idx + 1, + pool=pool, + memory_engine=memory_engine, + llm_config=llm_config, + request_context=request_context, + config=config, + ) - stats["observations_deleted"] += all_deleted - results = all_results + batch_results = await asyncio.gather( + *(_bounded_batch(i, b) for i, b in enumerate(llm_batches)) + ) + llm_batch_num += len(llm_batches) - # Checkpoint: abort if the operation (and thus the bank) was deleted mid-run. - if operation_id and not await memory_engine._check_op_alive(operation_id): - logger.info( - f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early" + # Aggregate succeeded/failed IDs across all batches into a single DB + # commit per fetch round. Each batch's IDs are disjoint, so this is + # functionally equivalent to the previous per-batch commits. + all_succeeded_ids: list[Any] = [mid for br in batch_results for mid in br.succeeded_ids] + all_failed_ids: list[Any] = [mid for br in batch_results for mid in br.failed_ids] + async with acquire_with_retry(pool) as conn: + if all_succeeded_ids: + await conn.executemany( + f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1", + [(mem_id,) for mem_id in all_succeeded_ids], + ) + if all_failed_ids: + await conn.executemany( + f"UPDATE {fq_table('memory_units')} SET consolidation_failed_at = NOW() WHERE id = $1", + [(mem_id,) for mem_id in all_failed_ids], ) - return {"status": "cancelled", "bank_id": bank_id, **stats} - for result in results: + # Merge per-batch perf into the shared parent perf, accumulate tags + stats. + # Per-batch perf was captured into a local ConsolidationPerfLog inside each + # task so timings/llm_calls aren't raced across concurrent gather participants. + for br in batch_results: + perf.merge(br.perf) + consolidated_tags.update(br.tags) + stats["observations_deleted"] += br.deleted + for result in br.results: stats["memories_processed"] += 1 action = result.get("action") if action == "created": @@ -544,29 +644,51 @@ async def run_consolidation_job( elif action == "failed": stats["memories_failed"] += 1 - # Per-LLM-batch log - llm_batch_time = time.time() - llm_batch_start - timing_parts = [] - for key in ["recall", "llm", "embedding", "db_write"]: - if key in perf.timings: - delta = perf.timings[key] - snap_timings.get(key, 0) - timing_parts.append(f"{key}={delta:.3f}s") - input_tokens = int((perf.total_prompt_chars - snap_total_chars) / 4) - batch_created = stats["observations_created"] - snap_stats["observations_created"] - batch_updated = stats["observations_updated"] - snap_stats["observations_updated"] - batch_skipped = stats["skipped"] - snap_stats["skipped"] - batch_failed = stats["memories_failed"] - snap_stats["memories_failed"] - llm_calls_made = perf.llm_calls - snap_llm_calls + # Per-batch logs (ordered by batch_num — accurate per-batch timings + # from each batch's own perf snapshot, no cross-batch contamination + # of the kind the previous snap-delta pattern would have produced + # under parallel execution). + for br in batch_results: + timing_parts = [ + f"{k}={br.perf.timings[k]:.3f}s" + for k in ("recall", "llm", "embedding", "db_write") + if k in br.perf.timings + ] + input_tokens = int(br.perf.total_prompt_chars / 4) + batch_created = sum( + (1 if r.get("action") == "created" else 0) + + (r.get("created", 0) if r.get("action") == "multiple" else 0) + for r in br.results + ) + batch_updated = sum( + (1 if r.get("action") == "updated" else 0) + + (r.get("updated", 0) if r.get("action") == "multiple" else 0) + for r in br.results + ) + batch_skipped = sum(1 for r in br.results if r.get("action") == "skipped") + batch_failed = sum(1 for r in br.results if r.get("action") == "failed") + denom = max(1, br.memories_count) logger.info( - f"[CONSOLIDATION] bank={bank_id} llm_batch #{llm_batch_num}" - f" ({len(llm_batch)} memories, {llm_calls_made} llm calls)" + f"[CONSOLIDATION] bank={bank_id} llm_batch #{br.batch_num}" + f" ({br.memories_count} memories, {br.perf.llm_calls} llm calls)" f" | {stats['memories_processed']}/{total_count} processed" f" | {', '.join(timing_parts)}" f" | created={batch_created} updated={batch_updated} skipped={batch_skipped}" + (f" failed={batch_failed}" if batch_failed else "") + f" | input_tokens=~{input_tokens}" - f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory" + f" | avg={br.elapsed / denom:.3f}s/memory" + ) + + # Checkpoint: abort if the operation (and thus the bank) was deleted + # during this gather wave. Granularity is per-wave instead of per-batch + # in the previous design; acceptable since waves are bounded by + # consolidation_batch_size memories (default 50) and waves typically + # complete in seconds once thinking is off. + if operation_id and not await memory_engine._check_op_alive(operation_id): + logger.info( + f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early" ) + return {"status": "cancelled", "bank_id": bank_id, **stats} # Update round budget after processing this DB fetch batch if round_limit_enabled: From 150c7de87c55938a69585f281347b19730d662b5 Mon Sep 17 00:00:00 2001 From: Connor Black Date: Tue, 12 May 2026 21:25:39 -0400 Subject: [PATCH 2/4] refactor(consolidation): tighten parallelism patch per upstream review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address reviewer feedback on the prior commit before opening upstream PR: - ConsolidationPerfLog.merge() -> __iadd__: matches the codebase's TokenUsage.__add__ accumulator idiom (used heavily in fact_extraction). Callers now write `perf += batch_result.perf`. - Extract _resolve_obs_tags_list() so observation_scopes parsing happens once per llm_batch instead of once per sub_batch (all sub_batches share the same parent batch's tags by tag-grouping invariant). - Extract _apply_action_to_stats() so the action-vocabulary mapping has one definition; per-batch counters and aggregate stats now come from a single pass over batch_result.results instead of two. - Plumb operation_id through _execute_one_llm_batch with a per-sub-batch cancellation check via _check_op_alive — restores per-batch cancellation granularity that the prior commit traded for per-wave only. - Tighten three docstrings (merge, _BatchExecutionResult, _execute_one_llm_batch) to contracts; drop refactor-narrating paragraphs. - Inline _resolve_obs_tags_list and the shared-tags assignment at the top of _execute_one_llm_batch, replacing the per-sub-batch redundant computation and the for-memory tag-tracking loop (all memories in a batch share tags). - Comment in run_consolidation_job notes the new semaphore stacks on top of the global LLM semaphore in llm_wrapper.py — effective concurrency is min(this, the global cap). - Rename loop variable b -> batch and br -> batch_result for readability. - Drop dead `denom = max(1, br.memories_count)` guard (memories_count is always >= 1 for non-empty batches by construction); use `or 1` inline. No behavior changes intended. Smoke-tested end-to-end on a deployed image. --- .../engine/consolidation/consolidator.py | 274 ++++++++---------- 1 file changed, 127 insertions(+), 147 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index 261bd1e68c..319b27e3b0 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -216,20 +216,14 @@ def record_llm_call(self, obs_count: int, prompt_chars: int) -> None: self.total_obs_in_context += obs_count self.total_prompt_chars += prompt_chars - def merge(self, other: "ConsolidationPerfLog") -> None: - """Merge another perf log's contributions into this one. - - Used to safely aggregate per-batch perf state collected by concurrent - consolidation tasks (see _execute_one_llm_batch). Each task gets its own - ConsolidationPerfLog so timings/llm_calls/prompt_chars are not racing - on the shared parent perf during asyncio.gather, then merged in-order - after gather completes. - """ + def __iadd__(self, other: "ConsolidationPerfLog") -> "ConsolidationPerfLog": + """Add another log's timings, llm_calls, and char counts into this one.""" for key, val in other.timings.items(): self.timings[key] = self.timings.get(key, 0) + val self.llm_calls += other.llm_calls self.total_obs_in_context += other.total_obs_in_context self.total_prompt_chars += other.total_prompt_chars + return self def flush(self) -> None: """Flush all log lines to the logger.""" @@ -243,13 +237,7 @@ def flush(self) -> None: @dataclass class _BatchExecutionResult: - """Self-contained result of one llm_batch execution. - - Returned by ``_execute_one_llm_batch`` so the parent ``run_consolidation_job`` - can aggregate stats, perf, and DB writes from concurrent batch tasks after - asyncio.gather completes. Each field is independent — no shared mutable state - is touched during batch execution; aggregation is single-threaded after gather. - """ + """Per-llm_batch outputs for parent aggregation after asyncio.gather.""" batch_num: int memories_count: int @@ -262,6 +250,62 @@ class _BatchExecutionResult: elapsed: float +def _apply_action_to_stats(stats: dict[str, int], result: dict[str, Any]) -> dict[str, int]: + """Apply one consolidation result to ``stats`` and return per-batch counters. + + Returns a small dict of ``{created, updated, skipped, failed}`` deltas so + callers can both update aggregate stats and emit per-batch log lines from a + single pass over results. + """ + counters = {"created": 0, "updated": 0, "skipped": 0, "failed": 0} + stats["memories_processed"] += 1 + action = result.get("action") + if action == "created": + stats["observations_created"] += 1 + stats["actions_executed"] += 1 + counters["created"] = 1 + elif action == "updated": + stats["observations_updated"] += 1 + stats["actions_executed"] += 1 + counters["updated"] = 1 + elif action == "merged": + stats["observations_merged"] += 1 + stats["actions_executed"] += 1 + elif action == "multiple": + c = result.get("created", 0) + u = result.get("updated", 0) + stats["observations_created"] += c + stats["observations_updated"] += u + stats["observations_merged"] += result.get("merged", 0) + stats["actions_executed"] += result.get("total_actions", 0) + counters["created"] = c + counters["updated"] = u + elif action == "skipped": + stats["skipped"] += 1 + counters["skipped"] = 1 + elif action == "failed": + stats["memories_failed"] += 1 + counters["failed"] = 1 + return counters + + +def _resolve_obs_tags_list(memory_tags: list[str], scope_spec: Any) -> list[list[str]] | None: + """Translate a memory's ``observation_scopes`` value into concrete tag-set passes. + + Returns ``None`` for single-pass (combined) consolidation, or a list of tag + sets for multi-pass consolidation (one observation pass per returned tag set). + """ + if scope_spec == "per_tag": + return [[tag] for tag in memory_tags] if memory_tags else None + if scope_spec == "all_combinations": + if not memory_tags: + return None + return [list(combo) for r in range(1, len(memory_tags) + 1) for combo in combinations(memory_tags, r)] + if scope_spec == "combined" or scope_spec is None: + return None + return scope_spec # explicit list[list[str]] + + async def _execute_one_llm_batch( *, bank_id: str, @@ -272,18 +316,16 @@ async def _execute_one_llm_batch( llm_config: Any, request_context: "RequestContext", config: Any, + operation_id: str | None = None, ) -> _BatchExecutionResult: - """Execute one llm_batch with adaptive split-on-failure, returning a self-contained result. + """Execute one llm_batch with adaptive split-on-failure. - Designed to be safely run concurrently with other batches via asyncio.gather - under a Semaphore at config.consolidation_llm_max_concurrent. All mutable - state (perf, succeeded_ids, failed_ids, tags) is local to this invocation; - nothing is shared with other concurrent invocations. The caller merges - these results into shared state after gather completes. + All mutable state is local — safe to run concurrently via asyncio.gather. + Adaptive split (halve sub-batch on LLM failure, mark failed only at size=1) + remains serial within this batch; ordering is required by the split protocol. - The adaptive split logic (halve sub-batch on LLM failure, retry, mark failed - only at sub_batch=1) remains serial within this batch — that ordering is a - correctness requirement of the split protocol. + ``operation_id`` enables per-sub-batch cancellation polling so a deleted bank + is honored mid-batch instead of waiting for the whole gather wave to drain. """ batch_start = time.time() batch_perf = ConsolidationPerfLog(bank_id) @@ -292,50 +334,26 @@ async def _execute_one_llm_batch( all_deleted = 0 succeeded_ids: list[Any] = [] failed_ids: list[Any] = [] - batch_tags: set[str] = set() - - # Track tags for mental model refresh filtering (per-batch local set) - for memory in llm_batch: - memory_tags = memory.get("tags") or [] - if memory_tags: - batch_tags.update(memory_tags) - - # Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch - # and retry, down to batch_size=1. Only if a single-memory batch still fails is - # the memory marked with consolidation_failed_at and excluded from future runs - # until explicitly retried via the API. + + # All memories in one llm_batch share the same tag set (enforced by tag_groups upstream), + # so observation_scopes is identical too — parse once for the whole batch. + shared_tags: list[str] = list(llm_batch[0].get("tags") or []) if llm_batch else [] + batch_tags: set[str] = set(shared_tags) + _obs_raw = llm_batch[0].get("observation_scopes") if llm_batch else None + _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw + obs_tags_list = _resolve_obs_tags_list(shared_tags, _obs_parsed) + pending: list[list[dict[str, Any]]] = [llm_batch] while pending: + # Per-sub-batch cancellation: a deleted bank is honored mid-wave rather + # than waiting for the whole gather to drain. + if operation_id and not await memory_engine._check_op_alive(operation_id): + break + sub_batch = pending.pop(0) + # One conn per sub_batch (multi-pass consolidation reuses it across passes). async with acquire_with_retry(pool) as conn: - # Determine observation_scopes for this sub-batch. All memories share - # the same tags (enforced by tag_groups), so we only check the first memory. - # asyncpg returns JSONB columns as raw JSON strings, so parse if needed. - _obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None - _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw - - # Resolve the scope spec into a concrete list[list[str]] (or None for combined). - if _obs_parsed == "per_tag": - _memory_tags = sub_batch[0].get("tags") or [] - obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None - elif _obs_parsed == "all_combinations": - _memory_tags = sub_batch[0].get("tags") or [] - obs_tags_list = ( - [ - list(combo) - for r in range(1, len(_memory_tags) + 1) - for combo in combinations(_memory_tags, r) - ] - if _memory_tags - else None - ) - elif _obs_parsed == "combined" or _obs_parsed is None: - obs_tags_list = None # single combined pass (default behaviour) - else: - # explicit list[list[str]] - obs_tags_list = _obs_parsed - sub_deleted: int = 0 sub_llm_failed = False if obs_tags_list: @@ -566,43 +584,38 @@ async def run_consolidation_job( for i in range(0, len(group), llm_batch_size): llm_batches.append(group[i : i + llm_batch_size]) - # Parallel batch execution with semaphore at consolidation_llm_max_concurrent. - # Each llm_batch operates on disjoint memories (the outer SELECT … LIMIT - # query returns a fixed set this round; tag-group slicing produces - # non-overlapping subsets), so concurrent execution is safe — no two - # batches can touch the same memory_id. The tag-group security boundary - # is preserved unchanged: cross-tag memories never share an LLM call. + # Each llm_batch operates on disjoint memories (tag-group slicing produces + # non-overlapping subsets), so concurrent execution cannot race on memory_id. + # The tag-group security boundary is preserved: cross-tag memories never + # share an LLM call. # - # Pre-refactor (serial for-loop) left consolidation_llm_max_concurrent - # unused for single-bank workloads, bottlenecking on single-LLM-call - # latency. With asyncio.gather + Semaphore we saturate up to N parallel - # LLM calls within one consolidation op for an N× throughput gain. + # The bound here stacks on top of the global LLM semaphore in + # llm_wrapper.py — effective concurrency is min(this, the global cap). sem = asyncio.Semaphore(max(1, config.consolidation_llm_max_concurrent)) - batch_base_num = llm_batch_num + start_num = llm_batch_num # captured because llm_batch_num is reassigned below - async def _bounded_batch(idx: int, b: list[dict[str, Any]]) -> _BatchExecutionResult: + async def _bounded_batch(idx: int, batch: list[dict[str, Any]]) -> _BatchExecutionResult: async with sem: return await _execute_one_llm_batch( bank_id=bank_id, - llm_batch=b, - llm_batch_num=batch_base_num + idx + 1, + llm_batch=batch, + llm_batch_num=start_num + idx + 1, pool=pool, memory_engine=memory_engine, llm_config=llm_config, request_context=request_context, config=config, + operation_id=operation_id, ) - batch_results = await asyncio.gather( - *(_bounded_batch(i, b) for i, b in enumerate(llm_batches)) - ) + batch_results = await asyncio.gather(*(_bounded_batch(i, batch) for i, batch in enumerate(llm_batches))) llm_batch_num += len(llm_batches) - # Aggregate succeeded/failed IDs across all batches into a single DB - # commit per fetch round. Each batch's IDs are disjoint, so this is - # functionally equivalent to the previous per-batch commits. - all_succeeded_ids: list[Any] = [mid for br in batch_results for mid in br.succeeded_ids] - all_failed_ids: list[Any] = [mid for br in batch_results for mid in br.failed_ids] + # Single DB commit for all succeeded/failed IDs across the wave. + # Each batch's IDs are disjoint by tag-grouping, so this matches the + # previous per-batch commits' effect with one round-trip instead of N. + all_succeeded_ids: list[Any] = [mid for batch_result in batch_results for mid in batch_result.succeeded_ids] + all_failed_ids: list[Any] = [mid for batch_result in batch_results for mid in batch_result.failed_ids] async with acquire_with_retry(pool) as conn: if all_succeeded_ids: await conn.executemany( @@ -615,75 +628,42 @@ async def _bounded_batch(idx: int, b: list[dict[str, Any]]) -> _BatchExecutionRe [(mem_id,) for mem_id in all_failed_ids], ) - # Merge per-batch perf into the shared parent perf, accumulate tags + stats. - # Per-batch perf was captured into a local ConsolidationPerfLog inside each - # task so timings/llm_calls aren't raced across concurrent gather participants. - for br in batch_results: - perf.merge(br.perf) - consolidated_tags.update(br.tags) - stats["observations_deleted"] += br.deleted - for result in br.results: - stats["memories_processed"] += 1 - action = result.get("action") - if action == "created": - stats["observations_created"] += 1 - stats["actions_executed"] += 1 - elif action == "updated": - stats["observations_updated"] += 1 - stats["actions_executed"] += 1 - elif action == "merged": - stats["observations_merged"] += 1 - stats["actions_executed"] += 1 - elif action == "multiple": - stats["observations_created"] += result.get("created", 0) - stats["observations_updated"] += result.get("updated", 0) - stats["observations_merged"] += result.get("merged", 0) - stats["actions_executed"] += result.get("total_actions", 0) - elif action == "skipped": - stats["skipped"] += 1 - elif action == "failed": - stats["memories_failed"] += 1 - - # Per-batch logs (ordered by batch_num — accurate per-batch timings - # from each batch's own perf snapshot, no cross-batch contamination - # of the kind the previous snap-delta pattern would have produced - # under parallel execution). - for br in batch_results: + # Single pass: merge per-batch perf, accumulate tags + stats, and emit + # the per-batch log line. Each batch's perf was captured locally so + # timings/llm_calls don't race across concurrent gather participants. + for batch_result in batch_results: + perf += batch_result.perf + consolidated_tags.update(batch_result.tags) + stats["observations_deleted"] += batch_result.deleted + + batch_counters = {"created": 0, "updated": 0, "skipped": 0, "failed": 0} + for result in batch_result.results: + deltas = _apply_action_to_stats(stats, result) + for k, v in deltas.items(): + batch_counters[k] += v + timing_parts = [ - f"{k}={br.perf.timings[k]:.3f}s" + f"{k}={batch_result.perf.timings[k]:.3f}s" for k in ("recall", "llm", "embedding", "db_write") - if k in br.perf.timings + if k in batch_result.perf.timings ] - input_tokens = int(br.perf.total_prompt_chars / 4) - batch_created = sum( - (1 if r.get("action") == "created" else 0) - + (r.get("created", 0) if r.get("action") == "multiple" else 0) - for r in br.results - ) - batch_updated = sum( - (1 if r.get("action") == "updated" else 0) - + (r.get("updated", 0) if r.get("action") == "multiple" else 0) - for r in br.results - ) - batch_skipped = sum(1 for r in br.results if r.get("action") == "skipped") - batch_failed = sum(1 for r in br.results if r.get("action") == "failed") - denom = max(1, br.memories_count) + input_tokens = int(batch_result.perf.total_prompt_chars / 4) + failed_part = f" failed={batch_counters['failed']}" if batch_counters["failed"] else "" logger.info( - f"[CONSOLIDATION] bank={bank_id} llm_batch #{br.batch_num}" - f" ({br.memories_count} memories, {br.perf.llm_calls} llm calls)" + f"[CONSOLIDATION] bank={bank_id} llm_batch #{batch_result.batch_num}" + f" ({batch_result.memories_count} memories, {batch_result.perf.llm_calls} llm calls)" f" | {stats['memories_processed']}/{total_count} processed" f" | {', '.join(timing_parts)}" - f" | created={batch_created} updated={batch_updated} skipped={batch_skipped}" - + (f" failed={batch_failed}" if batch_failed else "") - + f" | input_tokens=~{input_tokens}" - f" | avg={br.elapsed / denom:.3f}s/memory" + f" | created={batch_counters['created']}" + f" updated={batch_counters['updated']}" + f" skipped={batch_counters['skipped']}" + f"{failed_part}" + f" | input_tokens=~{input_tokens}" + f" | avg={batch_result.elapsed / (batch_result.memories_count or 1):.3f}s/memory" ) - # Checkpoint: abort if the operation (and thus the bank) was deleted - # during this gather wave. Granularity is per-wave instead of per-batch - # in the previous design; acceptable since waves are bounded by - # consolidation_batch_size memories (default 50) and waves typically - # complete in seconds once thinking is off. + # End-of-wave checkpoint (per-sub-batch checks inside _execute_one_llm_batch + # already shorten the in-wave cancellation window). if operation_id and not await memory_engine._check_op_alive(operation_id): logger.info( f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early" From d21faa1605ad21b8b59047bd85efbf77348d05b2 Mon Sep 17 00:00:00 2001 From: Connor Black Date: Tue, 12 May 2026 21:33:41 -0400 Subject: [PATCH 3/4] refactor(consolidation): round-2 simplifications per upstream review Address remaining findings before opening upstream PR: - Introduce _ResultDelta dataclass with __iadd__; _record_result returns _ResultDelta instead of an opaque dict. Caller now writes `batch_counters += _record_result(stats, result)` (matches the codebase's TokenUsage.__add__ accumulator idiom). - Track 'merged' in per-batch counters and per-batch log line. Pre-patch silently dropped merged actions from the per-batch log; the new helper is the right place to fix this. - Extract _merge_pass_result(existing, new) -> dict from the dense inline block in _execute_one_llm_batch's multi-pass loop. Centralizes the skipped-is-weak / non-skipped-combine-into-multiple action vocabulary that previously duplicated knowledge across two sites. - Extract _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; the implicit `operation_id is not None` short-circuit is now in one place. - Replace `start_num = llm_batch_num` capture pattern with `enumerate(llm_batches, start=llm_batch_num + 1)` to match the codebase's idiom (memory_engine.py:3234, search/tracer.py:325, etc.). - Drop dead `(memories_count or 1)` defensive guard. memories_count is always >= 1 by tag-group construction (range slice over non-empty group). - Rename _apply_action_to_stats -> _record_result. Old name suggested one-way write; new name covers the mutate-and-return shape clearly. --- .../engine/consolidation/consolidator.py | 141 +++++++++++------- 1 file changed, 86 insertions(+), 55 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index 319b27e3b0..fb991d22d0 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -250,43 +250,97 @@ class _BatchExecutionResult: elapsed: float -def _apply_action_to_stats(stats: dict[str, int], result: dict[str, Any]) -> dict[str, int]: - """Apply one consolidation result to ``stats`` and return per-batch counters. +@dataclass +class _ResultDelta: + """Per-result counters returned by ``_record_result`` for per-batch log lines.""" + + created: int = 0 + updated: int = 0 + merged: int = 0 + skipped: int = 0 + failed: int = 0 + + def __iadd__(self, other: "_ResultDelta") -> "_ResultDelta": + self.created += other.created + self.updated += other.updated + self.merged += other.merged + self.skipped += other.skipped + self.failed += other.failed + return self + - Returns a small dict of ``{created, updated, skipped, failed}`` deltas so - callers can both update aggregate stats and emit per-batch log lines from a - single pass over results. +def _record_result(stats: dict[str, int], result: dict[str, Any]) -> _ResultDelta: + """Apply one consolidation result to aggregate ``stats`` and return per-result deltas. + + Both mutates ``stats`` and returns a delta so callers can build per-batch + counters in a single pass over results. """ - counters = {"created": 0, "updated": 0, "skipped": 0, "failed": 0} + delta = _ResultDelta() stats["memories_processed"] += 1 action = result.get("action") if action == "created": stats["observations_created"] += 1 stats["actions_executed"] += 1 - counters["created"] = 1 + delta.created = 1 elif action == "updated": stats["observations_updated"] += 1 stats["actions_executed"] += 1 - counters["updated"] = 1 + delta.updated = 1 elif action == "merged": stats["observations_merged"] += 1 stats["actions_executed"] += 1 + delta.merged = 1 elif action == "multiple": c = result.get("created", 0) u = result.get("updated", 0) + m = result.get("merged", 0) stats["observations_created"] += c stats["observations_updated"] += u - stats["observations_merged"] += result.get("merged", 0) + stats["observations_merged"] += m stats["actions_executed"] += result.get("total_actions", 0) - counters["created"] = c - counters["updated"] = u + delta.created = c + delta.updated = u + delta.merged = m elif action == "skipped": stats["skipped"] += 1 - counters["skipped"] = 1 + delta.skipped = 1 elif action == "failed": stats["memories_failed"] += 1 - counters["failed"] = 1 - return counters + delta.failed = 1 + return delta + + +def _merge_pass_result(existing: dict[str, Any], new: dict[str, Any]) -> dict[str, Any]: + """Combine two per-memory results from sibling multi-pass consolidation passes. + + ``skipped`` is the weak default — any non-skipped pass wins. Two non-skipped + passes combine into ``"multiple"`` with summed created/updated counts. + """ + existing_action = existing.get("action") + new_action = new.get("action") + if existing_action == "skipped": + return new + if new_action == "skipped": + return existing + + existing_created = existing.get("created", 1 if existing_action == "created" else 0) + existing_updated = existing.get("updated", 1 if existing_action == "updated" else 0) + new_created = new.get("created", 1 if new_action == "created" else 0) + new_updated = new.get("updated", 1 if new_action == "updated" else 0) + return { + "action": "multiple", + "created": existing_created + new_created, + "updated": existing_updated + new_updated, + "merged": 0, + "total_actions": existing_created + existing_updated + new_created + new_updated, + } + + +async def _is_op_cancelled(memory_engine: "MemoryEngine", operation_id: str | None) -> bool: + """True when ``operation_id`` is set and its async_operations row is no longer alive.""" + if not operation_id: + return False + return not await memory_engine._check_op_alive(operation_id) def _resolve_obs_tags_list(memory_tags: list[str], scope_spec: Any) -> list[list[str]] | None: @@ -345,9 +399,8 @@ async def _execute_one_llm_batch( pending: list[list[dict[str, Any]]] = [llm_batch] while pending: - # Per-sub-batch cancellation: a deleted bank is honored mid-wave rather - # than waiting for the whole gather to drain. - if operation_id and not await memory_engine._check_op_alive(operation_id): + # Honor a deleted bank mid-wave instead of waiting for gather to drain. + if await _is_op_cancelled(memory_engine, operation_id): break sub_batch = pending.pop(0) @@ -357,7 +410,6 @@ async def _execute_one_llm_batch( sub_deleted: int = 0 sub_llm_failed = False if obs_tags_list: - # Multi-pass: run one observation consolidation pass per tag set sub_results: list[dict[str, Any]] = [] for obs_tags in obs_tags_list: pass_results, pass_deleted, pass_failed = await _process_memory_batch( @@ -373,33 +425,11 @@ async def _execute_one_llm_batch( ) sub_deleted += pass_deleted sub_llm_failed = sub_llm_failed or pass_failed - # Merge results: prefer non-skipped actions if not sub_results: sub_results = pass_results else: - for i, (existing, new) in enumerate(zip(sub_results, pass_results)): - if existing.get("action") == "skipped" and new.get("action") != "skipped": - sub_results[i] = new - elif existing.get("action") != "skipped" and new.get("action") != "skipped": - # Both did something — combine into "multiple" - existing_created = existing.get( - "created", 1 if existing.get("action") == "created" else 0 - ) - existing_updated = existing.get( - "updated", 1 if existing.get("action") == "updated" else 0 - ) - new_created = new.get("created", 1 if new.get("action") == "created" else 0) - new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0) - total = existing_created + existing_updated + new_created + new_updated - sub_results[i] = { - "action": "multiple", - "created": existing_created + new_created, - "updated": existing_updated + new_updated, - "merged": 0, - "total_actions": total, - } + sub_results = [_merge_pass_result(e, n) for e, n in zip(sub_results, pass_results)] else: - # Normal single pass using the memory's own tags sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch( conn=conn, memory_engine=memory_engine, @@ -592,14 +622,13 @@ async def run_consolidation_job( # The bound here stacks on top of the global LLM semaphore in # llm_wrapper.py — effective concurrency is min(this, the global cap). sem = asyncio.Semaphore(max(1, config.consolidation_llm_max_concurrent)) - start_num = llm_batch_num # captured because llm_batch_num is reassigned below - async def _bounded_batch(idx: int, batch: list[dict[str, Any]]) -> _BatchExecutionResult: + async def _bounded_batch(batch_num: int, batch: list[dict[str, Any]]) -> _BatchExecutionResult: async with sem: return await _execute_one_llm_batch( bank_id=bank_id, llm_batch=batch, - llm_batch_num=start_num + idx + 1, + llm_batch_num=batch_num, pool=pool, memory_engine=memory_engine, llm_config=llm_config, @@ -608,7 +637,9 @@ async def _bounded_batch(idx: int, batch: list[dict[str, Any]]) -> _BatchExecuti operation_id=operation_id, ) - batch_results = await asyncio.gather(*(_bounded_batch(i, batch) for i, batch in enumerate(llm_batches))) + batch_results = await asyncio.gather( + *(_bounded_batch(num, batch) for num, batch in enumerate(llm_batches, start=llm_batch_num + 1)) + ) llm_batch_num += len(llm_batches) # Single DB commit for all succeeded/failed IDs across the wave. @@ -636,11 +667,9 @@ async def _bounded_batch(idx: int, batch: list[dict[str, Any]]) -> _BatchExecuti consolidated_tags.update(batch_result.tags) stats["observations_deleted"] += batch_result.deleted - batch_counters = {"created": 0, "updated": 0, "skipped": 0, "failed": 0} + batch_counters = _ResultDelta() for result in batch_result.results: - deltas = _apply_action_to_stats(stats, result) - for k, v in deltas.items(): - batch_counters[k] += v + batch_counters += _record_result(stats, result) timing_parts = [ f"{k}={batch_result.perf.timings[k]:.3f}s" @@ -648,23 +677,25 @@ async def _bounded_batch(idx: int, batch: list[dict[str, Any]]) -> _BatchExecuti if k in batch_result.perf.timings ] input_tokens = int(batch_result.perf.total_prompt_chars / 4) - failed_part = f" failed={batch_counters['failed']}" if batch_counters["failed"] else "" + failed_part = f" failed={batch_counters.failed}" if batch_counters.failed else "" + merged_part = f" merged={batch_counters.merged}" if batch_counters.merged else "" logger.info( f"[CONSOLIDATION] bank={bank_id} llm_batch #{batch_result.batch_num}" f" ({batch_result.memories_count} memories, {batch_result.perf.llm_calls} llm calls)" f" | {stats['memories_processed']}/{total_count} processed" f" | {', '.join(timing_parts)}" - f" | created={batch_counters['created']}" - f" updated={batch_counters['updated']}" - f" skipped={batch_counters['skipped']}" + f" | created={batch_counters.created}" + f" updated={batch_counters.updated}" + f"{merged_part}" + f" skipped={batch_counters.skipped}" f"{failed_part}" f" | input_tokens=~{input_tokens}" - f" | avg={batch_result.elapsed / (batch_result.memories_count or 1):.3f}s/memory" + f" | avg={batch_result.elapsed / batch_result.memories_count:.3f}s/memory" ) # End-of-wave checkpoint (per-sub-batch checks inside _execute_one_llm_batch # already shorten the in-wave cancellation window). - if operation_id and not await memory_engine._check_op_alive(operation_id): + if await _is_op_cancelled(memory_engine, operation_id): logger.info( f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early" ) From 1fbe9971d0b4bffbe3b7d1a3a18e76a6fd34a626 Mon Sep 17 00:00:00 2001 From: Connor Black Date: Tue, 12 May 2026 22:26:12 -0400 Subject: [PATCH 4/4] fix(consolidation): address upstream review on parallelism PR Address Copilot review on #1604: 1. Coalesce None on consolidation_llm_max_concurrent The field is Optional[int] in HindsightConfig and only set when HINDSIGHT_API_CONSOLIDATION_LLM_MAX_CONCURRENT is in env. Without the coalesce, max(1, None) raises TypeError on first wave. Fall back to the global llm_max_concurrent cap (which always has a default). 2. Use return_exceptions=True on the gather wave Without it, gather's first-exception-cancels-siblings semantic would skip the post-wave UPDATE, leaving observation rows already inserted by successful batches without their consolidated_at marker. Those memories would be re-consolidated on the next run, producing duplicate observations. Now: partition gather results into successes vs the first exception, apply DB markers for successful batches, then re-raise so the worker poller's standard exception handling kicks in. Additional exceptions in the same wave are logged via exc_info. 3. Revert obs_tags caching in _execute_one_llm_batch Round-2 hoisted observation_scopes parsing out of the while-pending loop on the assumption tag_groups upstream guaranteed uniform scopes per batch. tag_groups actually only keys on tags, so adaptive split sub_batches could legitimately have different observation_scopes than the parent llm_batch. Parse per sub_batch to preserve scope semantics. 4. Add tests/test_consolidation_parallelism.py - test_consolidation_honors_max_concurrent: pin max=4, ingest 8 memories with distinct tag groups, mock _process_memory_batch with a slot counter; assert 1 < max_in_flight <= 4. - test_partial_failure_preserves_succeeded_markers: ingest 6 memories, mock raises on the 2nd call; assert 5 succeeded markers present, 1 absent, exception re-raised. --- .../engine/consolidation/consolidator.py | 58 ++++++-- .../tests/test_consolidation_parallelism.py | 137 ++++++++++++++++++ 2 files changed, 182 insertions(+), 13 deletions(-) create mode 100644 hindsight-api-slim/tests/test_consolidation_parallelism.py diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index fb991d22d0..deded6716d 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -389,13 +389,12 @@ async def _execute_one_llm_batch( succeeded_ids: list[Any] = [] failed_ids: list[Any] = [] - # All memories in one llm_batch share the same tag set (enforced by tag_groups upstream), - # so observation_scopes is identical too — parse once for the whole batch. + # tag_groups upstream keys only on tags, not on observation_scopes — so + # observation_scopes can vary across sub_batches even though tags don't. + # Resolve obs_tags_list per sub_batch from sub_batch[0] to preserve + # per-memory scope semantics under adaptive split. shared_tags: list[str] = list(llm_batch[0].get("tags") or []) if llm_batch else [] batch_tags: set[str] = set(shared_tags) - _obs_raw = llm_batch[0].get("observation_scopes") if llm_batch else None - _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw - obs_tags_list = _resolve_obs_tags_list(shared_tags, _obs_parsed) pending: list[list[dict[str, Any]]] = [llm_batch] while pending: @@ -404,6 +403,9 @@ async def _execute_one_llm_batch( break sub_batch = pending.pop(0) + _obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None + _obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw + obs_tags_list = _resolve_obs_tags_list(sub_batch[0].get("tags") or [] if sub_batch else [], _obs_parsed) # One conn per sub_batch (multi-pass consolidation reuses it across passes). async with acquire_with_retry(pool) as conn: @@ -619,9 +621,12 @@ async def run_consolidation_job( # The tag-group security boundary is preserved: cross-tag memories never # share an LLM call. # - # The bound here stacks on top of the global LLM semaphore in - # llm_wrapper.py — effective concurrency is min(this, the global cap). - sem = asyncio.Semaphore(max(1, config.consolidation_llm_max_concurrent)) + # consolidation_llm_max_concurrent is Optional[int] in HindsightConfig; + # fall back to the global llm_max_concurrent cap when unset. The bound + # here stacks on top of the global LLM semaphore in llm_wrapper.py — + # effective concurrency is min(this, the global cap). + consolidation_concurrency = config.consolidation_llm_max_concurrent or config.llm_max_concurrent + sem = asyncio.Semaphore(max(1, consolidation_concurrency)) async def _bounded_batch(batch_num: int, batch: list[dict[str, Any]]) -> _BatchExecutionResult: async with sem: @@ -637,14 +642,35 @@ async def _bounded_batch(batch_num: int, batch: list[dict[str, Any]]) -> _BatchE operation_id=operation_id, ) - batch_results = await asyncio.gather( - *(_bounded_batch(num, batch) for num, batch in enumerate(llm_batches, start=llm_batch_num + 1)) + # return_exceptions=True so a single batch failure does not cancel siblings + # mid-wave. Without this, gather's first-exception cancellation would skip + # the post-wave UPDATE, leaving observation rows already inserted by + # successful batches without their consolidated_at marker — those memories + # would be re-consolidated on the next run, producing duplicate observations. + gather_results = await asyncio.gather( + *(_bounded_batch(num, batch) for num, batch in enumerate(llm_batches, start=llm_batch_num + 1)), + return_exceptions=True, ) llm_batch_num += len(llm_batches) - # Single DB commit for all succeeded/failed IDs across the wave. - # Each batch's IDs are disjoint by tag-grouping, so this matches the - # previous per-batch commits' effect with one round-trip instead of N. + batch_results: list[_BatchExecutionResult] = [] + first_batch_exc: BaseException | None = None + for r in gather_results: + if isinstance(r, BaseException): + if first_batch_exc is None: + first_batch_exc = r + else: + logger.error( + f"[CONSOLIDATION] bank={bank_id} additional batch exception suppressed: {r!r}", + exc_info=r, + ) + else: + batch_results.append(r) + + # Single DB commit for all succeeded/failed IDs across the wave's + # successful batches. Each batch's IDs are disjoint by tag-grouping; + # IDs from cancelled/raised batches are intentionally absent so those + # memories will be re-attempted on the next consolidation run. all_succeeded_ids: list[Any] = [mid for batch_result in batch_results for mid in batch_result.succeeded_ids] all_failed_ids: list[Any] = [mid for batch_result in batch_results for mid in batch_result.failed_ids] async with acquire_with_retry(pool) as conn: @@ -693,6 +719,12 @@ async def _bounded_batch(batch_num: int, batch: list[dict[str, Any]]) -> _BatchE f" | avg={batch_result.elapsed / batch_result.memories_count:.3f}s/memory" ) + # Re-raise the first batch exception now that DB state is consistent + # for everything that did succeed. The worker poller's standard + # exception handling marks the operation row failed for retry. + if first_batch_exc is not None: + raise first_batch_exc + # End-of-wave checkpoint (per-sub-batch checks inside _execute_one_llm_batch # already shorten the in-wave cancellation window). if await _is_op_cancelled(memory_engine, operation_id): diff --git a/hindsight-api-slim/tests/test_consolidation_parallelism.py b/hindsight-api-slim/tests/test_consolidation_parallelism.py new file mode 100644 index 0000000000..7094f1ed0a --- /dev/null +++ b/hindsight-api-slim/tests/test_consolidation_parallelism.py @@ -0,0 +1,137 @@ +"""Tests for the bounded-parallel consolidation path. + +Covers: +- ``consolidation_llm_max_concurrent`` is honored (waves see N in-flight calls). +- Exception in one batch does not strand successful batches' consolidated_at + markers (gather uses ``return_exceptions=True``; succeeded IDs commit before + the exception is re-raised). +""" + +import asyncio +import uuid +from unittest.mock import AsyncMock, patch + +import pytest + +from hindsight_api.config import _get_raw_config +from hindsight_api.engine.consolidation.consolidator import run_consolidation_job +from hindsight_api.engine.memory_engine import MemoryEngine + + +@pytest.fixture +def consolidation_concurrency(): + """Pin observations + consolidation_llm_max_concurrent for deterministic concurrency assertions.""" + config = _get_raw_config() + original_obs = config.enable_observations + original_concurrency = config.consolidation_llm_max_concurrent + config.enable_observations = True + config.consolidation_llm_max_concurrent = 4 + try: + yield 4 + finally: + config.enable_observations = original_obs + config.consolidation_llm_max_concurrent = original_concurrency + + +async def _ingest_distinct_tag_groups(memory, request_context, bank_id, count): + """Retain ``count`` memories with distinct single-tag groups so each becomes its own llm_batch.""" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + for i in range(count): + await memory.retain_async( + bank_id=bank_id, + content=f"Person{i} prefers activity{i}.", + tags=[f"person{i}"], + request_context=request_context, + ) + + +class TestConsolidationParallelism: + """Bounded-parallel execution of llm_batches within one consolidation op.""" + + @pytest.mark.asyncio + async def test_consolidation_honors_max_concurrent( + self, memory: MemoryEngine, request_context, consolidation_concurrency + ): + """A wave of N llm_batches yields ``N <= max_concurrent`` simultaneously + in-flight LLM calls — proving the semaphore + gather pair are wired.""" + bank_id = f"test-parallelism-{uuid.uuid4().hex[:8]}" + n_memories = 8 # > max_concurrent (4) so the cap is exercised + + await _ingest_distinct_tag_groups(memory, request_context, bank_id, n_memories) + + in_flight = 0 + max_in_flight = 0 + gate = asyncio.Lock() + + async def slow_mock_process(*args, **kwargs): + nonlocal in_flight, max_in_flight + async with gate: + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + try: + # Hold the slot long enough for siblings to enter + await asyncio.sleep(0.05) + finally: + async with gate: + in_flight -= 1 + mems = kwargs.get("memories") or args[4] + return ([{"action": "skipped"} for _ in mems], 0, False) + + with patch( + "hindsight_api.engine.consolidation.consolidator._process_memory_batch", + new=AsyncMock(side_effect=slow_mock_process), + ): + await run_consolidation_job(memory, bank_id, request_context) + + assert max_in_flight > 1, f"expected concurrent batches, saw max_in_flight={max_in_flight}" + assert max_in_flight <= consolidation_concurrency, ( + f"max_in_flight={max_in_flight} exceeded cap {consolidation_concurrency}" + ) + + @pytest.mark.asyncio + async def test_partial_failure_preserves_succeeded_markers( + self, memory: MemoryEngine, request_context, consolidation_concurrency + ): + """When one batch raises, sibling batches that succeeded must still get + consolidated_at = NOW(); the exception is re-raised after DB consistency.""" + bank_id = f"test-parallelism-fail-{uuid.uuid4().hex[:8]}" + n_memories = 6 + + await _ingest_distinct_tag_groups(memory, request_context, bank_id, n_memories) + + # Make the second LLM call raise; siblings succeed. + call_index = 0 + call_index_lock = asyncio.Lock() + + async def mock_process(*args, **kwargs): + nonlocal call_index + async with call_index_lock: + call_index += 1 + idx = call_index + if idx == 2: + raise RuntimeError("synthetic LLM failure for batch 2") + mems = kwargs.get("memories") or args[4] + return ([{"action": "created"} for _ in mems], 0, False) + + with patch( + "hindsight_api.engine.consolidation.consolidator._process_memory_batch", + new=AsyncMock(side_effect=mock_process), + ), pytest.raises(RuntimeError, match="synthetic LLM failure"): + await run_consolidation_job(memory, bank_id, request_context) + + # Successful batches' memories must have consolidated_at set; the failed + # batch's memory must not (will be retried on next consolidation run). + async with memory._backend.acquire() as conn: + rows = await conn.fetch( + "SELECT id, consolidated_at FROM memory_units " + "WHERE bank_id = $1 AND fact_type IN ('experience', 'world')", + bank_id, + ) + consolidated = sum(1 for r in rows if r["consolidated_at"] is not None) + unconsolidated = sum(1 for r in rows if r["consolidated_at"] is None) + assert consolidated == n_memories - 1, ( + f"expected {n_memories - 1} consolidated (sibling success), got {consolidated}" + ) + assert unconsolidated == 1, ( + f"expected 1 unconsolidated (failed batch), got {unconsolidated}" + )