From 241ce6aecd7096246c03e561303c2cd182867549 Mon Sep 17 00:00:00 2001 From: oskar Date: Wed, 5 Aug 2026 12:02:41 +0200 Subject: [PATCH] feat(kb-mail-batching): retry + izolacja trujacego chunka w torze embed + benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Batching /api/embed juz istnial (Krok 1 fazy mailowej, batch 64). Recon przed Etapem B wykazal w torze backfillu blad blokujacy i dwie luki. BUG (blokujacy dla Etapu B): flush_embed_buffer lapal wylacznie aiohttp.ClientError, a wyczerpanie ClientTimeout(total=...) rzuca goly builtins.TimeoutError, ktory NIE jest jego podklasa (zweryfikowane empirycznie na aiohttp 3.14.3). Zawieszona Ollama — czyli jej udokumentowany failure mode, "przyjmuje polaczenie i milczy" — wywalala caly run nieobsluzonym wyjatkiem, bez breakera i bez flushu threadingu. Na plastrze 50k = utrata zarobionej pracy. Klasy przejsciowe nazwane teraz jawnie w TRANSIENT_EMBED_ERRORS. kb-retrieval: - embed_batch(timeout_s=...) — bound per zadanie, skalowalny z batch size - embed_batch_resilient() — retry z backoffem wykladniczym, a po ich wyczerpaniu probe /api/tags rozstrzyga: backend zywy -> bisekcja izolujaca trujacy chunk (jeden zly tekst kosztowal caly batch 64, bo /api/embed jest all-or-nothing); backend martwy -> natychmiastowe gave_up bez bisekcji, ktora spalilaby 2n-1 zadan i opoznila breaker. EmbeddingDimensionError nigdy nie jest retry'owane. - failed_indices wyprowadzane z wyniku, nie akumulowane per span — przy gave_up w srodku bisekcji porzucone poddrzewo nigdy nie dochodzi do liscia. mail-body-ingest: - breaker liczy give-upy (backend padl), nie dowolne nieudane batche; porazka czesciowa przy zywym backendzie nie przesuwa licznika, bo te chunki i tak zlapie kolejny run przez idempotencje - wiersze zembedowane w umierajacym batchu sa commitowane przed abortem - parametryzacja: --batch-size/--embed-retries/--embed-backoff/--embed-timeout, kazdy z odpowiednikiem env MAIL_INGEST_*; bledna wartosc env = glosny SystemExit - metryka embed_ms_per_chunk (porownywalna miedzy runami, w odroznieniu od sredniej per batch) + embed_requests_total/embed_calls jako sygnal zdrowia mail-body-ingest-bench: nowy entry point, sweep batch size na realnych chunkach. Read-only (SELECT + inferencja, zero sciezki zapisu), warmup przed pomiarem, ten sam zbior chunkow dla kazdego rozmiaru. Czyni liczby z planu §1.4 odtwarzalnymi. Fallback SOLARIA->PIHA dla backfillu SWIADOMIE nie powstaje (potwierdzone przez operatora): 271k chunkow x 790 ms CPU ~ 60 h na 8 GB PIHA dzielonym z HA i Paperlessem. Wlasciwa odpowiedzia na martwy backend jest exit 2 i wznowienie plastra. Tor online (kb-query -> embed_router) zachowuje fallback — rozdzial torow udokumentowany w docstringu embed.py i w kb/services/. Testy: 117 zielonych (62 job + 22 klient embed + reszta pakietow), w tym regresja na TimeoutError, bisekcja, ograniczony koszt przy martwym backendzie i porazka czesciowa nieprzesuwajaca breakera. Bez uruchamiania backfillu. Co-Authored-By: Claude Opus 5 (1M context) --- jobs/mail-body-ingest/pyproject.toml | 1 + .../src/mail_body_ingest/benchmark.py | 241 ++++++++++++++++++ .../src/mail_body_ingest/ingest.py | 184 ++++++++++--- jobs/mail-body-ingest/tests/test_benchmark.py | 155 +++++++++++ jobs/mail-body-ingest/tests/test_ingest.py | 122 ++++++++- kb/phases/kb-m5-faza-mailowa.md | 31 ++- kb/runbooks/mail-body-ingest-run.md | 52 +++- kb/services/job-mail-body-ingest.md | 81 ++++-- .../kb-retrieval/src/kb_retrieval/embed.py | 160 +++++++++++- packages/kb-retrieval/tests/test_embed.py | 198 +++++++++++++- 10 files changed, 1143 insertions(+), 82 deletions(-) create mode 100644 jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py create mode 100644 jobs/mail-body-ingest/tests/test_benchmark.py diff --git a/jobs/mail-body-ingest/pyproject.toml b/jobs/mail-body-ingest/pyproject.toml index 9ed1ead..8163950 100644 --- a/jobs/mail-body-ingest/pyproject.toml +++ b/jobs/mail-body-ingest/pyproject.toml @@ -22,6 +22,7 @@ dev = [ [project.scripts] mail-body-ingest = "mail_body_ingest.ingest:main" +mail-body-ingest-bench = "mail_body_ingest.benchmark:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py b/jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py new file mode 100644 index 0000000..abb4b95 --- /dev/null +++ b/jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py @@ -0,0 +1,241 @@ +"""Batch-size benchmark for the mail backfill's embed path. + +A separate entry point rather than a `--benchmark` flag on the job, deliberately: this module +imports nothing that writes. It issues `SELECT`s and Ollama inference and has no code path that +can INSERT, UPDATE or touch a container — so it is safe to point at the live kb-postgres@PIHA +while the corpus it measures is the real one. + +The numbers in `kb/phases/kb-m5-faza-mailowa.md` §1.4 (batch 64 -> ~8-18 ms/chunk vs ~207 ms +sequential) were measured ad hoc during recon and are not reproducible. This makes them a +command, so batch sizing can be re-checked after an Ollama upgrade, on a different GPU, or when +a slice of the corpus turns out to have an unusual length distribution. + +Measures the SAME chunk set at every batch size — the comparison is otherwise meaningless, since +ms/chunk depends heavily on text length (§1.4: 8.5 ms/szt at ~900 chars vs 18.1 at ~3500). + +Usage (from SOLARIA, where Ollama and the .eml mirror live): + mail-body-ingest-bench --dsn postgresql://kb:@piha:5433/kb \ + --archive-root /home/oskar/kb/mail/archive --sample-envelopes 200 + + # Custom sweep, capped so batch=1 doesn't dominate the wall clock: + mail-body-ingest-bench --dsn ... --sizes 1,8,32,64,128 --max-chunks 300 +""" +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from dataclasses import dataclass +from pathlib import Path + +import aiohttp +import asyncpg +import structlog +from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text +from kb_retrieval.embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, embed_batch + +from mail_body_ingest.ingest import ( + DEFAULT_ARCHIVE_ROOT, + _decode_jsonb, + build_prefix, + extract_body, + fetch_envelopes, + is_newsletter, + parse_message, + strip_quotes, +) + +_log = structlog.get_logger(__name__) + +DEFAULT_SIZES = (1, 8, 32, 64) +DEFAULT_SAMPLE_ENVELOPES = 200 +# Active (non-newsletter) chunks projected for the full corpus — plan §1.3's extrapolation. +# Only used to turn ms/chunk into an "Etap B would take N hours" column. +FULL_CORPUS_CHUNKS = 271_000 + + +def split_batches(items: list, size: int) -> list[list]: + """Split into consecutive batches of `size`; the final batch may be short.""" + if size < 1: + raise ValueError(f"batch size must be >= 1, got {size}") + return [items[i : i + size] for i in range(0, len(items), size)] + + +@dataclass +class BenchRow: + batch_size: int + chunks: int + requests: int + seconds: float + failures: int = 0 + + @property + def ms_per_chunk(self) -> float: + return 1000 * self.seconds / self.chunks if self.chunks else 0.0 + + @property + def chunks_per_s(self) -> float: + return self.chunks / self.seconds if self.seconds else 0.0 + + @property + def projected_hours(self) -> float: + """Wall-clock hours to embed the full active corpus at this rate — GPU time only, + excluding the single-threaded parse the real run interleaves with it.""" + return FULL_CORPUS_CHUNKS * self.ms_per_chunk / 1000 / 3600 + + +def format_table(rows: list[BenchRow]) -> str: + header = ( + f"{'batch':>6} | {'chunks':>7} | {'requests':>8} | {'total s':>8} | " + f"{'ms/chunk':>9} | {'chunks/s':>9} | {'271k in h':>10} | {'fail':>5}" + ) + lines = [header, "-" * len(header)] + for r in rows: + lines.append( + f"{r.batch_size:>6} | {r.chunks:>7} | {r.requests:>8} | {r.seconds:>8.2f} | " + f"{r.ms_per_chunk:>9.2f} | {r.chunks_per_s:>9.1f} | {r.projected_hours:>10.2f} | " + f"{r.failures:>5}" + ) + if rows: + best = min(rows, key=lambda r: r.ms_per_chunk if r.chunks else float("inf")) + baseline = next((r for r in rows if r.batch_size == 1), None) + note = f"\nbest: batch={best.batch_size} at {best.ms_per_chunk:.2f} ms/chunk" + if baseline and baseline.ms_per_chunk and best is not baseline: + note += f" ({baseline.ms_per_chunk / best.ms_per_chunk:.1f}x faster than batch=1)" + lines.append(note) + return "\n".join(lines) + + +async def collect_chunks( + dsn: str, archive_root: Path, sample_envelopes: int, max_chunks: int | None +) -> list[str]: + """Read-only: the same parse -> quote-strip -> prefix -> chunk pipeline the job runs, so the + measured texts are the real workload. Newsletters are skipped because the real run never + embeds them (they go in with embedding=NULL), and mails whose file is missing or unparseable + are skipped silently — this is a benchmark, not a corpus audit.""" + conn = await asyncpg.connect(dsn) + try: + envelopes = await fetch_envelopes(conn, None, sample_envelopes, None) + finally: + await conn.close() + + chunks: list[str] = [] + for row in envelopes: + try: + raw = (archive_root / row["raw_ref"]).read_bytes() + msg = parse_message(raw) + if is_newsletter(msg): + continue + clean_body, _ = strip_quotes(extract_body(msg)) + except Exception: + continue + if not clean_body.strip(): + continue + entities = _decode_jsonb(row["entities"]) or [] + prefix = build_prefix(entities, row["ts"]) + chunks.extend(chunk_text(f"{prefix}\n\n{clean_body}", TARGET_CHARS, OVERLAP_CHARS)) + if max_chunks is not None and len(chunks) >= max_chunks: + return chunks[:max_chunks] + return chunks + + +async def bench_size( + session: aiohttp.ClientSession, ollama_url: str, model: str, chunks: list[str], + batch_size: int, timeout_s: float, +) -> BenchRow: + """Time `chunks` through `/api/embed` at one batch size. Failures are counted, not raised — + a size that the backend can't handle (OOM at batch 128, say) is itself a result, and must + not abort the remaining sizes in the sweep.""" + row = BenchRow(batch_size=batch_size, chunks=0, requests=0, seconds=0.0) + for batch in split_batches(chunks, batch_size): + try: + _embeddings, elapsed = await embed_batch( + session, ollama_url, model, batch, timeout_s=timeout_s + ) + except Exception as exc: + _log.warning("bench.batch_failed", batch_size=batch_size, error=str(exc)) + row.failures += 1 + continue + row.requests += 1 + row.chunks += len(batch) + row.seconds += elapsed + return row + + +async def run_benchmark( + dsn: str, archive_root: Path, ollama_url: str, model: str, sizes: list[int], + sample_envelopes: int, max_chunks: int | None, timeout_s: float, +) -> list[BenchRow]: + chunks = await collect_chunks(dsn, archive_root, sample_envelopes, max_chunks) + if not chunks: + raise SystemExit("no chunks collected — check --archive-root and --sample-envelopes") + _log.info("bench.corpus", chunks=len(chunks), + avg_chars=round(sum(len(c) for c in chunks) / len(chunks))) + + rows: list[BenchRow] = [] + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout_s)) as session: + # Warm-up: the first call after an idle period pays the model load (ollama-piha even runs + # OLLAMA_KEEP_ALIVE=0). Charging that to whichever size happens to run first would make + # the sweep unreproducible and unfairly penalise batch=1. + _log.info("bench.warmup", model=model, ollama_url=ollama_url) + await embed_batch(session, ollama_url, model, chunks[:1], timeout_s=timeout_s) + + for size in sizes: + _log.info("bench.size_start", batch_size=size, chunks=len(chunks)) + row = await bench_size(session, ollama_url, model, chunks, size, timeout_s) + _log.info("bench.size_done", batch_size=size, ms_per_chunk=round(row.ms_per_chunk, 2)) + rows.append(row) + return rows + + +def _parse_sizes(value: str) -> list[int]: + sizes = [int(v) for v in value.split(",") if v.strip()] + if not sizes or any(s < 1 for s in sizes): + raise argparse.ArgumentTypeError(f"invalid --sizes: {value!r}") + return sizes + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Benchmark Ollama /api/embed batch sizes on real mail chunks. " + "Read-only: SELECTs and inference, never writes." + ) + parser.add_argument("--dsn", default=os.environ.get("KB_DSN"), + help="asyncpg DSN for kb-postgres (or set KB_DSN env var)") + parser.add_argument("--archive-root", type=Path, default=DEFAULT_ARCHIVE_ROOT, + help=f"Mail .eml archive root (default: {DEFAULT_ARCHIVE_ROOT})") + parser.add_argument("--ollama-url", default=os.environ.get("OLLAMA_URL", DEFAULT_OLLAMA_URL), + help=f"Ollama base URL (default: {DEFAULT_OLLAMA_URL})") + parser.add_argument("--model", default=os.environ.get("OLLAMA_EMBED_MODEL", DEFAULT_MODEL), + help=f"Ollama embedding model (default: {DEFAULT_MODEL})") + parser.add_argument("--sizes", type=_parse_sizes, default=list(DEFAULT_SIZES), + help="Comma-separated batch sizes to sweep " + f"(default: {','.join(str(s) for s in DEFAULT_SIZES)})") + parser.add_argument("--sample-envelopes", type=int, default=DEFAULT_SAMPLE_ENVELOPES, + help=f"Envelopes to draw chunks from (default: {DEFAULT_SAMPLE_ENVELOPES})") + parser.add_argument("--max-chunks", type=int, default=None, + help="Cap the measured chunk set — batch=1 dominates the wall clock " + "otherwise (default: no cap)") + parser.add_argument("--timeout", type=float, default=300.0, metavar="SECONDS", + help="Per-request timeout (default: 300)") + args = parser.parse_args() + + if not args.dsn: + _log.error("missing_dsn", hint="pass --dsn or set KB_DSN") + sys.exit(1) + if not args.archive_root.is_dir(): + _log.error("archive_root_not_found", path=str(args.archive_root)) + sys.exit(1) + + rows = asyncio.run(run_benchmark( + dsn=args.dsn, archive_root=args.archive_root, ollama_url=args.ollama_url, + model=args.model, sizes=args.sizes, sample_envelopes=args.sample_envelopes, + max_chunks=args.max_chunks, timeout_s=args.timeout, + )) + print() + print(format_table(rows)) + + +if __name__ == "__main__": + main() diff --git a/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py b/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py index a9e9c4b..852df65 100644 --- a/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py +++ b/jobs/mail-body-ingest/src/mail_body_ingest/ingest.py @@ -42,17 +42,31 @@ Newsletter chunks (`List-Unsubscribe`/`List-Id`/`Precedence: bulk|list`) are ins stays in the DB (reversible: `UPDATE ... SET excluded_reason=NULL WHERE excluded_reason= 'newsletter'` + a re-embed run un-flags them later, per plan Decyzja 4). -Ollama-offline tolerance: a failed `embed_batch()` call is caught per-batch (`chunks_errors += -len(batch)`) — those chunks never enter the idempotency set, so a later re-run naturally retries -them. Transient single-batch failures therefore cost nothing. What they must NOT do is let a run -survive a *dead* backend: with the archive parsed single-threaded ahead of the GPU, a full-corpus -run (Etap B, plan §9) would otherwise keep chewing through 200k+ mails at parse speed, marking -every chunk `chunks_errors`, and the whole multi-hour pass would have to be repeated. Hence the -circuit breaker: `--max-embed-failures` consecutive failed batches (default 5, `0` disables) -raise `EmbedBackendUnavailableError` and stop the run early with exit code 2, after flushing the -threading updates already earned. A single successful batch resets the counter. Only a wrong -embedding dimension is more severe (`EmbeddingDimensionError`, exit 1) — never silently indexes -a mismatched vector. +Ollama-offline tolerance, in three layers (`kb_retrieval.embed.embed_batch_resilient` implements +the first two): + +1. **Transient blip** — each embed batch is retried `--embed-retries` times with exponential + backoff (`--embed-backoff`). Costs seconds, loses nothing. +2. **One poison chunk** — `/api/embed` is all-or-nothing, so a single pathological text fails + the whole 64-chunk batch it lands in. When retries are exhausted but `/api/tags` says the + backend is alive, the batch is bisected until the bad inputs are isolated; the rest embed + normally and only the bad ones count as `chunks_errors`. They never enter the idempotency + set, so a re-run retries them. +3. **Dead backend** — the probe says down, so the batch "gives up" without bisecting. With the + archive parsed single-threaded ahead of the GPU, a full-corpus run (Etap B, plan §9) would + otherwise chew through 200k+ mails at parse speed marking every chunk `chunks_errors`, and + the whole multi-hour pass would have to be repeated. Hence the circuit breaker: + `--max-embed-failures` consecutive give-ups (default 5, `0` disables) raise + `EmbedBackendUnavailableError` and stop the run with exit code 2, after committing the rows + and threading updates already earned. Any successful batch resets the counter. + +Only a partial failure against a *live* backend deliberately does NOT advance the breaker — a +handful of bad chunks must not abort a 50k slice. A wrong embedding dimension remains the most +severe outcome (`EmbeddingDimensionError`, exit 1) — never silently indexes a mismatched vector. + +No SOLARIA->PIHA fallback here, by decision: PIHA's ~790 ms/embed CPU over ~271k chunks is ~60 h +on the infra node. See `kb_retrieval.embed`'s module docstring for the full argument — the two +backend paths (online query vs backfill) are separate on purpose. """ from __future__ import annotations @@ -80,17 +94,21 @@ from kb_retrieval.embed import ( EmbeddingDimensionError, _vector_literal, check_ollama_health, - embed_batch, + embed_batch_resilient, ) _log = structlog.get_logger(__name__) DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive") DEFAULT_BATCH_SIZE = 64 +DEFAULT_EMBED_RETRIES = 2 +DEFAULT_EMBED_BACKOFF_S = 1.0 +DEFAULT_EMBED_TIMEOUT_S = 120.0 THREADING_UPDATE_BATCH_SIZE = 500 -# Circuit breaker: consecutive failed embed batches that mean "the backend is down, not flaky". -# 5 x --batch-size chunks written off before stopping; Ollama@SOLARIA's known failure mode is -# total (container vanishes / network-detached), not partial, so this trips within seconds of it. +# Circuit breaker: consecutive embed batches that gave up, i.e. whose /api/tags probe said the +# backend is down rather than merely choking on one input. 5 x --batch-size chunks written off +# before stopping; Ollama@SOLARIA's known failure mode is total (container vanishes / +# network-detached), not partial, so this trips within seconds of it. DEFAULT_MAX_EMBED_FAILURES = 5 EXIT_EMBED_BACKEND_UNAVAILABLE = 2 @@ -376,6 +394,15 @@ def _rows_affected(command_tag: str) -> int: return int(command_tag.rsplit(" ", 1)[-1]) +def _embed_ms_per_chunk(stats: dict) -> float: + """Wall-clock ms of Ollama time per successfully embedded chunk — the figure directly + comparable with `mail-body-ingest-bench`'s table, unlike a per-batch average (which moves + with --batch-size and so can't be compared across runs).""" + if not stats["embed_texts_ok"]: + return 0.0 + return round(1000 * stats["embed_seconds_total"] / stats["embed_texts_ok"], 2) + + def _new_stats() -> dict: return { "mails_scanned": 0, @@ -391,8 +418,15 @@ def _new_stats() -> dict: "chunks_conflict_skipped": 0, "chunks_errors": 0, "quoted_chars_stripped_total": 0, + # Diagnostics — none of these participate in the balance equations. `embed_calls` counts + # flushed batches, `embed_requests_total` the HTTP requests they actually cost (retries + # and bisection included), so the ratio is the health signal: 1.0 = a clean run. "embed_calls": 0, - "embed_batch_failures": 0, # diagnostic only — not part of the balance equations + "embed_requests_total": 0, + "embed_retries_total": 0, + "embed_batch_failures": 0, + "embed_items_failed": 0, + "embed_texts_ok": 0, "embed_seconds_total": 0.0, "threading_updated": 0, "threading_already_present": 0, @@ -410,6 +444,9 @@ async def run( apply: bool = False, batch_size: int = DEFAULT_BATCH_SIZE, max_embed_failures: int = DEFAULT_MAX_EMBED_FAILURES, + embed_retries: int = DEFAULT_EMBED_RETRIES, + embed_backoff_s: float = DEFAULT_EMBED_BACKOFF_S, + embed_timeout_s: float = DEFAULT_EMBED_TIMEOUT_S, ) -> dict: """Process one --limit/--offset (optionally --since-filtered) slice of `source='gmail'` envelopes. dry-run (apply=False): parse + quote-strip + classify + chunk + count, zero @@ -427,7 +464,9 @@ async def run( session: Optional[aiohttp.ClientSession] = None if apply: - session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) + # Session-wide ceiling; `embed_batch_resilient` additionally bounds each individual + # embed request at the same value, so one knob governs both. + session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=embed_timeout_s)) if not await check_ollama_health(session, ollama_url, 3.0): _log.warning("ollama_down_at_start", ollama_url=ollama_url) @@ -440,28 +479,41 @@ async def run( if not embed_buffer: return texts = [t for (_eid, _idx, t) in embed_buffer] - try: - embeddings, elapsed = await embed_batch(session, ollama_url, model, texts) - except EmbeddingDimensionError: - raise - except aiohttp.ClientError: - _log.warning("skip.embed_batch_error", count=len(embed_buffer), exc_info=True) - stats["chunks_errors"] += len(embed_buffer) - stats["embed_batch_failures"] += 1 - consecutive_embed_failures += 1 - embed_buffer.clear() - if max_embed_failures and consecutive_embed_failures >= max_embed_failures: - raise EmbedBackendUnavailableError( - f"{consecutive_embed_failures} consecutive embed batches failed " - f"({ollama_url}) — stopping instead of parsing the rest of the archive " - f"into chunks_errors; re-run to retry the missed chunks" - ) - return + # EmbeddingDimensionError / ValueError propagate untouched (abort-run). Everything + # transport-shaped is handled inside, including the timeouts that used to escape + # `except aiohttp.ClientError` and kill the run outright. + outcome = await embed_batch_resilient( + session, ollama_url, model, texts, + retries=embed_retries, backoff_s=embed_backoff_s, timeout_s=embed_timeout_s, + ) - consecutive_embed_failures = 0 stats["embed_calls"] += 1 - stats["embed_seconds_total"] += elapsed - for (eid, idx, chunk), embedding in zip(embed_buffer, embeddings): + stats["embed_requests_total"] += outcome.requests + stats["embed_retries_total"] += outcome.retries + stats["embed_seconds_total"] += outcome.elapsed_s + stats["embed_texts_ok"] += outcome.ok_count + if outcome.failed_indices: + stats["embed_batch_failures"] += 1 + stats["embed_items_failed"] += len(outcome.failed_indices) + _log.warning( + "skip.embed_items_failed", failed=len(outcome.failed_indices), + of=len(texts), gave_up=outcome.gave_up, ollama_url=ollama_url, + envelope_ids=[embed_buffer[i][0] for i in outcome.failed_indices[:10]], + ) + + # Only a give-up ("the backend stopped answering") advances the breaker. A partial + # failure against a live backend must not: those chunks were never inserted, so the + # next run retries them through the ordinary idempotency path, and letting a handful + # of poison chunks abort a 50k slice would be strictly worse than skipping them. + if outcome.gave_up: + consecutive_embed_failures += 1 + elif outcome.ok_count: + consecutive_embed_failures = 0 + + for (eid, idx, chunk), embedding in zip(embed_buffer, outcome.embeddings): + if embedding is None: + stats["chunks_errors"] += 1 + continue try: command_tag = await insert_chunk(conn, eid, idx, chunk, embedding, model) except Exception: @@ -476,6 +528,19 @@ async def run( stats["chunks_inserted"] += 1 embed_buffer.clear() + # Raised only after the successes above are committed — a batch that partially + # embedded before the backend died keeps the rows it earned. + if ( + outcome.gave_up + and max_embed_failures + and consecutive_embed_failures >= max_embed_failures + ): + raise EmbedBackendUnavailableError( + f"{consecutive_embed_failures} consecutive embed batches gave up " + f"({ollama_url}) — stopping instead of parsing the rest of the archive " + f"into chunks_errors; re-run to retry the missed chunks" + ) + async def flush_threading() -> None: if apply and threading_pending: await conn.executemany(_THREADING_UPDATE_SQL, threading_pending) @@ -566,7 +631,7 @@ async def run( await flush_embed_buffer() if stats["mails_scanned"] % 500 == 0: - _log.info("progress", **stats) + _log.info("progress", embed_ms_per_chunk=_embed_ms_per_chunk(stats), **stats) if apply: await flush_embed_buffer() @@ -602,6 +667,19 @@ def _parse_since(value: str) -> datetime: return datetime.strptime(value, "%Y-%m-%d").replace(tzinfo=timezone.utc) +def _env_num(name: str, default, cast): + """Env override for a numeric CLI default. A malformed value is a loud startup failure, not a + silent fallback to the default — a typo'd MAIL_INGEST_BATCH_SIZE must not quietly produce a + multi-hour run at the wrong batch size.""" + raw = os.environ.get(name) + if raw is None or raw == "": + return default + try: + return cast(raw) + except ValueError: + raise SystemExit(f"{name}={raw!r} is not a valid {cast.__name__}") + + def main() -> None: parser = argparse.ArgumentParser( description="Chunk + embed source='gmail' envelope body content into document_chunk " @@ -621,8 +699,25 @@ def main() -> None: help="Max envelopes to process (default: all)") parser.add_argument("--offset", type=int, default=0, help="Slice offset, ordered by envelope id (default: 0)") - parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE, - help=f"Ollama /api/embed batch size (default: {DEFAULT_BATCH_SIZE})") + parser.add_argument("--batch-size", type=int, + default=_env_num("MAIL_INGEST_BATCH_SIZE", DEFAULT_BATCH_SIZE, int), + help=f"Ollama /api/embed batch size (default: {DEFAULT_BATCH_SIZE}, " + f"or set MAIL_INGEST_BATCH_SIZE)") + parser.add_argument("--embed-retries", type=int, + default=_env_num("MAIL_INGEST_EMBED_RETRIES", DEFAULT_EMBED_RETRIES, int), + help=f"Retries per embed batch on transport errors, exponential backoff " + f"(default: {DEFAULT_EMBED_RETRIES}, or set MAIL_INGEST_EMBED_RETRIES). " + f"0 disables retrying.") + parser.add_argument("--embed-backoff", type=float, + default=_env_num("MAIL_INGEST_EMBED_BACKOFF", DEFAULT_EMBED_BACKOFF_S, float), + metavar="SECONDS", + help=f"Base backoff between embed retries, doubled each attempt " + f"(default: {DEFAULT_EMBED_BACKOFF_S}, or set MAIL_INGEST_EMBED_BACKOFF)") + parser.add_argument("--embed-timeout", type=float, + default=_env_num("MAIL_INGEST_EMBED_TIMEOUT", DEFAULT_EMBED_TIMEOUT_S, float), + metavar="SECONDS", + help=f"Hard timeout per embed request — scale it with --batch-size " + f"(default: {DEFAULT_EMBED_TIMEOUT_S}, or set MAIL_INGEST_EMBED_TIMEOUT)") parser.add_argument("--max-embed-failures", type=int, default=DEFAULT_MAX_EMBED_FAILURES, metavar="N", help=f"Abort (exit {EXIT_EMBED_BACKEND_UNAVAILABLE}) after N consecutive " @@ -654,6 +749,9 @@ def main() -> None: apply=args.apply, batch_size=args.batch_size, max_embed_failures=args.max_embed_failures, + embed_retries=args.embed_retries, + embed_backoff_s=args.embed_backoff, + embed_timeout_s=args.embed_timeout, ) ) except EmbeddingDimensionError as exc: @@ -669,7 +767,11 @@ def main() -> None: avg_embed_ms = ( 1000 * stats["embed_seconds_total"] / stats["embed_calls"] if stats["embed_calls"] else 0.0 ) - _log.info("summary", mode=mode, avg_embed_ms_per_batch=round(avg_embed_ms, 2), **stats) + _log.info( + "summary", mode=mode, batch_size=args.batch_size, + avg_embed_ms_per_batch=round(avg_embed_ms, 2), + embed_ms_per_chunk=_embed_ms_per_chunk(stats), **stats, + ) balanced = ( stats["mails_scanned"] == ( diff --git a/jobs/mail-body-ingest/tests/test_benchmark.py b/jobs/mail-body-ingest/tests/test_benchmark.py new file mode 100644 index 0000000..9bd4ac8 --- /dev/null +++ b/jobs/mail-body-ingest/tests/test_benchmark.py @@ -0,0 +1,155 @@ +"""Unit tests for the batch-size benchmark — no real HTTP, no real Ollama, no DB.""" +from __future__ import annotations + +import aiohttp +import pytest + +from mail_body_ingest.benchmark import ( + FULL_CORPUS_CHUNKS, + BenchRow, + bench_size, + format_table, + split_batches, +) + + +class TestSplitBatches: + def test_exact_division(self): + assert split_batches([1, 2, 3, 4], 2) == [[1, 2], [3, 4]] + + def test_trailing_short_batch_is_kept(self): + assert split_batches([1, 2, 3, 4, 5], 2) == [[1, 2], [3, 4], [5]] + + def test_size_one_is_one_item_per_batch(self): + assert split_batches([1, 2, 3], 1) == [[1], [2], [3]] + + def test_size_larger_than_input_yields_a_single_batch(self): + assert split_batches([1, 2], 64) == [[1, 2]] + + def test_empty_input_yields_no_batches(self): + assert split_batches([], 8) == [] + + def test_covers_every_item_exactly_once(self): + items = list(range(100)) + flattened = [i for batch in split_batches(items, 7) for i in batch] + assert flattened == items + + def test_zero_or_negative_size_raises(self): + with pytest.raises(ValueError): + split_batches([1, 2], 0) + + +class TestBenchRow: + def test_ms_per_chunk(self): + row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=2.0) + assert row.ms_per_chunk == 20.0 + + def test_chunks_per_second(self): + row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=2.0) + assert row.chunks_per_s == 50.0 + + def test_projected_hours_scales_the_measured_rate_to_the_full_corpus(self): + row = BenchRow(batch_size=64, chunks=100, requests=2, seconds=1.0) # 10 ms/chunk + assert row.projected_hours == pytest.approx(FULL_CORPUS_CHUNKS * 0.010 / 3600) + + def test_zero_chunks_does_not_divide_by_zero(self): + row = BenchRow(batch_size=1, chunks=0, requests=0, seconds=0.0) + assert row.ms_per_chunk == 0.0 + assert row.chunks_per_s == 0.0 + + +class TestFormatTable: + def _rows(self): + return [ + BenchRow(batch_size=1, chunks=100, requests=100, seconds=20.0), # 200 ms/chunk + BenchRow(batch_size=64, chunks=100, requests=2, seconds=1.0), # 10 ms/chunk + ] + + def test_one_line_per_row_plus_header(self): + out = format_table(self._rows()) + assert "batch" in out and "ms/chunk" in out + assert "200.00" in out and "10.00" in out + + def test_names_the_best_size_and_the_speedup_over_batch_1(self): + out = format_table(self._rows()) + assert "best: batch=64" in out + assert "20.0x faster than batch=1" in out + + def test_no_speedup_note_when_batch_1_wins(self): + rows = [BenchRow(batch_size=1, chunks=100, requests=100, seconds=1.0)] + out = format_table(rows) + assert "best: batch=1" in out + assert "faster than" not in out + + def test_empty_rows_still_renders_a_header(self): + assert "batch" in format_table([]) + + def test_a_size_that_failed_entirely_is_not_reported_as_best(self): + """A batch size the backend chokes on measures 0 chunks; ranking it best would read as + 'infinitely fast' instead of 'did not work'.""" + rows = [ + BenchRow(batch_size=8, chunks=100, requests=13, seconds=2.0), + BenchRow(batch_size=512, chunks=0, requests=0, seconds=0.0, failures=1), + ] + assert "best: batch=8" in format_table(rows) + + +class _FakeResponse: + def __init__(self, payload, fail=False): + self._payload = payload + self._fail = fail + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def json(self): + return self._payload + + def raise_for_status(self): + if self._fail: + raise aiohttp.ClientConnectionError("simulated failure") + + +class _FakeSession: + def __init__(self, dim=1024, fail_sizes=()): + self._dim = dim + self._fail_sizes = set(fail_sizes) + self.posts: list[int] = [] + + def post(self, url, json, timeout=None): + texts = json["input"] + self.posts.append(len(texts)) + if len(texts) in self._fail_sizes: + return _FakeResponse({}, fail=True) + return _FakeResponse({"embeddings": [[0.01] * self._dim for _ in texts]}) + + +class TestBenchSize: + async def test_measures_every_chunk_once(self): + session = _FakeSession() + chunks = [f"c{i}" for i in range(10)] + row = await bench_size(session, "http://fake", "bge-m3", chunks, 4, 60.0) + assert row.chunks == 10 + assert row.requests == 3 # 4 + 4 + 2 + assert session.posts == [4, 4, 2] + assert row.failures == 0 + + async def test_failures_are_counted_not_raised(self): + """A size the backend can't serve is a result, not a reason to abandon the sweep.""" + session = _FakeSession(fail_sizes={8}) + chunks = [f"c{i}" for i in range(16)] + row = await bench_size(session, "http://fake", "bge-m3", chunks, 8, 60.0) + assert row.failures == 2 + assert row.chunks == 0 + assert row.ms_per_chunk == 0.0 + + async def test_partial_failure_still_measures_the_successful_batches(self): + session = _FakeSession(fail_sizes={2}) # only the short trailing batch fails + chunks = [f"c{i}" for i in range(10)] + row = await bench_size(session, "http://fake", "bge-m3", chunks, 4, 60.0) + assert row.chunks == 8 + assert row.requests == 2 + assert row.failures == 1 diff --git a/jobs/mail-body-ingest/tests/test_ingest.py b/jobs/mail-body-ingest/tests/test_ingest.py index 4e74e2d..a85e987 100644 --- a/jobs/mail-body-ingest/tests/test_ingest.py +++ b/jobs/mail-body-ingest/tests/test_ingest.py @@ -19,7 +19,10 @@ from mail_body_ingest.ingest import ( _CHUNK_INSERT_SQL, EmbedBackendUnavailableError, _decode_jsonb, + _embed_ms_per_chunk, + _env_num, _has_threading, + _new_stats, build_prefix, extract_body, extract_threading, @@ -389,7 +392,13 @@ class _FakeTagsResponse: class _FakeOllamaSession: - def __init__(self, dim=1024, fail_batches=False, health_up=True, fail_pattern=None): + """`health_up` doubles as "is the backend alive": it answers the /api/tags probe that + `embed_batch_resilient` uses to tell a dead backend (-> give up, breaker counts it) from a + poison input on a live one (-> bisect, breaker unaffected). A test that wants the breaker + to trip must therefore set health_up=False, not merely fail the batches.""" + + def __init__(self, dim=1024, fail_batches=False, health_up=True, fail_pattern=None, + fail_texts=None): self._dim = dim self._fail_batches = fail_batches self._health_up = health_up @@ -397,23 +406,30 @@ class _FakeOllamaSession: # `fail_batches`. Lets a test interleave failures and successes to exercise the breaker's # "consecutive" semantics rather than a plain total. self._fail_pattern = list(fail_pattern) if fail_pattern is not None else None + # Substrings that poison whatever batch they land in — the all-or-nothing property of + # /api/embed that bisection exists to work around. + self._fail_texts = list(fail_texts or ()) self.requests: list[dict] = [] + self.health_checks = 0 self.closed = False - def _should_fail(self) -> bool: + def _should_fail(self, texts) -> bool: + if any(needle in t for t in texts for needle in self._fail_texts): + return True if self._fail_pattern: return self._fail_pattern.pop(0) return self._fail_batches - def post(self, url, json): + def post(self, url, json, timeout=None): assert url.endswith("/api/embed") self.requests.append({"url": url, "json": json}) - if self._should_fail(): + if self._should_fail(json["input"]): return _FakeEmbedResponse({}, status=500) n = len(json["input"]) return _FakeEmbedResponse({"embeddings": [[0.01] * self._dim for _ in range(n)]}) def get(self, url, timeout=None): + self.health_checks += 1 return _FakeTagsResponse(200 if self._health_up else 500) async def close(self): @@ -556,10 +572,11 @@ class TestRun: {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" )) conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")]) - ollama = _FakeOllamaSession(dim=1024, fail_batches=True) + ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False) self._patch(monkeypatch, conn, ollama) - stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True) + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, + embed_retries=0) assert stats["chunks_errors"] == 1 assert stats["chunks_inserted"] == 0 @@ -583,12 +600,12 @@ class TestRun: """A dead backend must stop the run, not let it parse the rest of the archive into chunks_errors (plan §9 Etap B: 200k+ mails behind a single-threaded parse).""" conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8)) - ollama = _FakeOllamaSession(dim=1024, fail_batches=True) + ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False) self._patch(monkeypatch, conn, ollama) with pytest.raises(EmbedBackendUnavailableError): await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, - batch_size=1, max_embed_failures=5) + batch_size=1, max_embed_failures=5, embed_retries=0) # Stopped at the 5th failed batch — mails 6-8 were never parsed, let alone embedded. assert len(ollama.requests) == 5 @@ -597,11 +614,12 @@ class TestRun: async def test_breaker_counter_resets_on_successful_batch(self, tmp_path, monkeypatch): """"Consecutive", not "total" — flaky batches interleaved with successes must not trip it.""" conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 5)) - ollama = _FakeOllamaSession(dim=1024, fail_pattern=[True, True, False, True, True]) + ollama = _FakeOllamaSession(dim=1024, health_up=False, + fail_pattern=[True, True, False, True, True]) self._patch(monkeypatch, conn, ollama) stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, - batch_size=1, max_embed_failures=3) + batch_size=1, max_embed_failures=3, embed_retries=0) assert len(ollama.requests) == 5 # ran to completion assert stats["chunks_errors"] == 4 @@ -614,11 +632,11 @@ class TestRun: async def test_breaker_disabled_with_zero(self, tmp_path, monkeypatch): conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 6)) - ollama = _FakeOllamaSession(dim=1024, fail_batches=True) + ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False) self._patch(monkeypatch, conn, ollama) stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, - batch_size=1, max_embed_failures=0) + batch_size=1, max_embed_failures=0, embed_retries=0) assert len(ollama.requests) == 6 assert stats["chunks_errors"] == 6 @@ -627,17 +645,93 @@ class TestRun: """Threading appends are Ollama-independent and idempotent — the abort keeps them rather than making the next run re-derive them from the same 27 GB read.""" conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8)) - ollama = _FakeOllamaSession(dim=1024, fail_batches=True) + ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False) self._patch(monkeypatch, conn, ollama) with pytest.raises(EmbedBackendUnavailableError): await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, - batch_size=1, max_embed_failures=5) + batch_size=1, max_embed_failures=5, embed_retries=0) assert len(conn.executemany_calls) == 1 # The 5 mails processed before the breaker tripped, none of the 3 after it. assert len(conn.executemany_calls[0][1]) == 5 + async def test_partial_batch_failure_does_not_advance_the_breaker(self, tmp_path, monkeypatch): + """A live backend that chokes on one chunk must not abort a 50k slice. The bad chunk is + isolated, everything else embeds, and the run continues past max_embed_failures batches.""" + rows = self._write_n_mails(tmp_path, 6) + # Mail 2's body poisons any batch it lands in; the backend stays healthy throughout. + ollama = _FakeOllamaSession(dim=1024, health_up=True, fail_texts=["number 2"]) + conn = _FakeConn(envelopes=rows) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, + batch_size=1, max_embed_failures=1, embed_retries=0) + + assert stats["chunks_errors"] == 1 # only the poison chunk + assert stats["chunks_inserted"] == 5 # the other five got through + assert stats["embed_items_failed"] == 1 + assert stats["mails_scanned"] == 6 # ran to completion, never aborted + assert stats["chunks_total"] == ( + stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + + stats["chunks_already_embedded"] + stats["chunks_conflict_skipped"] + + stats["chunks_errors"] + ) + + async def test_poison_chunk_isolated_from_a_full_batch(self, tmp_path, monkeypatch): + """Batched: one bad chunk used to cost the whole batch. Bisection keeps the rest.""" + conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8)) + ollama = _FakeOllamaSession(dim=1024, health_up=True, fail_texts=["number 5"]) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, + batch_size=8, embed_retries=0) + + assert stats["chunks_inserted"] == 7 + assert stats["chunks_errors"] == 1 + assert stats["embed_calls"] == 1 # one logical batch... + assert stats["embed_requests_total"] > 1 # ...several HTTP requests to isolate it + + async def test_embed_timeout_degrades_instead_of_crashing_the_run(self, tmp_path, monkeypatch): + """Regression: a bare builtins.TimeoutError from an exhausted ClientTimeout is not an + aiohttp.ClientError, so it used to escape the handler and kill the run — no breaker, no + threading flush. Ollama@SOLARIA hangs rather than refusing, so this was reachable.""" + conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 3)) + + class _HangingSession(_FakeOllamaSession): + def post(self, url, json, timeout=None): + self.requests.append({"url": url, "json": json}) + raise TimeoutError("simulated Ollama hang") + + ollama = _HangingSession(dim=1024, health_up=False) + self._patch(monkeypatch, conn, ollama) + + with pytest.raises(EmbedBackendUnavailableError): + await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True, + batch_size=1, max_embed_failures=2, embed_retries=0) + + # Reached the breaker (a clean stop) rather than propagating TimeoutError, and the + # threading work earned before the stop was still committed. + assert len(conn.executemany_calls) == 1 + + async def test_env_overrides_batch_size(self, monkeypatch): + monkeypatch.setenv("MAIL_INGEST_BATCH_SIZE", "16") + assert _env_num("MAIL_INGEST_BATCH_SIZE", 64, int) == 16 + + async def test_malformed_env_is_a_loud_failure(self, monkeypatch): + monkeypatch.setenv("MAIL_INGEST_BATCH_SIZE", "sixty-four") + with pytest.raises(SystemExit): + _env_num("MAIL_INGEST_BATCH_SIZE", 64, int) + + async def test_embed_ms_per_chunk_is_per_chunk_not_per_batch(self): + stats = _new_stats() + stats["embed_seconds_total"] = 2.0 + stats["embed_texts_ok"] = 100 + assert _embed_ms_per_chunk(stats) == 20.0 + + async def test_embed_ms_per_chunk_handles_zero_embeds(self): + assert _embed_ms_per_chunk(_new_stats()) == 0.0 + async def test_dimension_mismatch_aborts(self, tmp_path, monkeypatch): _write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml( {"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there" diff --git a/kb/phases/kb-m5-faza-mailowa.md b/kb/phases/kb-m5-faza-mailowa.md index 9f2d702..9d1ca9b 100644 --- a/kb/phases/kb-m5-faza-mailowa.md +++ b/kb/phases/kb-m5-faza-mailowa.md @@ -3,7 +3,7 @@ okf: "0.1" type: phase visibility: private status: active -updated: 2026-07-23 +updated: 2026-08-05 links: [] --- @@ -620,6 +620,35 @@ zatrzymaniu) wykazał dwie rzeczy do rozstrzygnięcia. Decyzje: 4. Przełączenie domyślnego `mode` kb-query na `hybrid` (DoD (d)) — **poza zakresem Etapu B**, osobny task po PASS regresji. +### Hardening toru embed przed Etapem B (2026-08-05) + +Recon pod kątem batchingu potwierdził, że batch `/api/embed` (Krok 1) działa zgodnie +z §1.4 — ale wykazał w torze backfillu **błąd blokujący dla Etapu B** i dwie luki: + +1. **BUG (naprawiony)**: `flush_embed_buffer` łapał wyłącznie `aiohttp.ClientError`, a + wyczerpanie `ClientTimeout(total=...)` rzuca goły `builtins.TimeoutError`, który **nie** + jest jego podklasą (zweryfikowane empirycznie na aiohttp 3.14.3). Zawieszona Ollama — + czyli dokładnie jej udokumentowany failure mode, „przyjmuje połączenie i milczy", nie + „odmawia" — wywalała cały run nieobsłużonym wyjątkiem: **bez breakera i bez flushu + threadingu**. Na plastrze 50k oznaczało to utratę już zarobionej pracy. Klasy przejściowe + nazwane teraz jawnie w `kb_retrieval.embed.TRANSIENT_EMBED_ERRORS`. +2. **Brak retry** — jeden blip sieciowy spisywał na straty cały batch (64 chunki). Dodane: + `--embed-retries` (default 2) z backoffem wykładniczym. +3. **Brak obsługi błędu częściowego** — `/api/embed` jest all-or-nothing, więc jeden trujący + chunk zabijał batch w kółko i mógł wywalić breaker przy **żywym** backendzie. Dodana + bisekcja po nieudanych retry, ale tylko gdy `/api/tags` potwierdza, że backend żyje; + przy martwym batch od razu „gives up" (bisekcja martwego backendu kosztowałaby 2n-1 + żądań i opóźniała breaker). Porażka częściowa **nie** przesuwa już breakera. + +Decyzja 2 (circuit breaker) obowiązuje w zaostrzonej formie: licznik liczy **give-upy** +(backend padł), nie dowolne nieudane batche. Fallback SOLARIA→PIHA dla backfillu **świadomie +nie powstaje** — 271k chunków × 790 ms CPU ≈ 60 h na 8 GB PIHA dzielonym z HA i Paperlessem; +właściwą odpowiedzią na martwy backend jest exit 2 i wznowienie plastra. Tor online (`kb-query` +→ `embed_router`) ma fallback i tak zostaje — te dwie ścieżki są rozdzielone celowo. + +Doszedł też `mail-body-ingest-bench` — sweep batch size na realnych chunkach (read-only), +żeby liczby z §1.4 dało się odtworzyć po zmianie GPU albo wersji Ollamy. + Uwaga do czytania wyników: na pełnym korpusie `exit 1` jest spodziewany (pojedyncze `parse_errors` — §1.5 dokumentuje ~9 maili na fallbacku compat32). Werdyktem jest bilans i liczniki w linii `summary`, nie kod wyjścia. `exit 2` diff --git a/kb/runbooks/mail-body-ingest-run.md b/kb/runbooks/mail-body-ingest-run.md index fed23b1..b6936c1 100644 --- a/kb/runbooks/mail-body-ingest-run.md +++ b/kb/runbooks/mail-body-ingest-run.md @@ -3,7 +3,7 @@ okf: "0.1" type: runbook visibility: private status: active -updated: 2026-07-22 +updated: 2026-08-05 links: - ../services/job-mail-body-ingest.md --- @@ -32,6 +32,34 @@ mail-body-ingest --dsn ... --apply --limit 10 DSN can also come from `KB_DSN`, Ollama URL from `OLLAMA_URL` (default `http://localhost:11434` — this job is meant to run where Ollama lives). +## Embed tuning + +| Flag | Env | Default | Notes | +|---|---|---|---| +| `--batch-size` | `MAIL_INGEST_BATCH_SIZE` | 64 | `/api/embed` batch. 64 measured best across the corpus's length mix; short-text-heavy slices favour 128 (plan §1.4) | +| `--embed-retries` | `MAIL_INGEST_EMBED_RETRIES` | 2 | Retries per batch on transport errors; `0` disables | +| `--embed-backoff` | `MAIL_INGEST_EMBED_BACKOFF` | 1.0 s | Base backoff, doubled each attempt | +| `--embed-timeout` | `MAIL_INGEST_EMBED_TIMEOUT` | 120 s | Per-request hard timeout — scale it with `--batch-size` | +| `--max-embed-failures` | — | 5 | Consecutive give-ups before aborting with exit 2; `0` disables | + +A malformed env value is a startup failure, not a silent fallback — a typo'd +`MAIL_INGEST_BATCH_SIZE` must not quietly produce a multi-hour run at the wrong batch size. + +## Batch-size benchmark + +Read-only (`SELECT`s + inference, no write path at all), so it is safe against the live DB: + +```bash +mail-body-ingest-bench --dsn postgresql://kb:@piha:5433/kb \ + --archive-root /home/oskar/kb/mail/archive --sample-envelopes 200 + +# Wider sweep, capped so batch=1 doesn't dominate the wall clock: +mail-body-ingest-bench --dsn ... --sizes 1,8,32,64,128 --max-chunks 300 +``` + +Prints ms/chunk, chunks/s and a projected full-corpus wall clock per batch size. Re-run it +after an Ollama upgrade or on a different GPU before trusting the default `--batch-size`. + ## Tests ```bash @@ -39,13 +67,27 @@ pip install -e "jobs/mail-body-ingest[dev]" cd jobs/mail-body-ingest && pytest ``` -Pure unit tests (55), no DB/Ollama — `run()` is tested by monkeypatching `asyncpg.connect` +Pure unit tests (62), no DB/Ollama — `run()` is tested by monkeypatching `asyncpg.connect` and `aiohttp.ClientSession` with in-memory fakes, `.eml` bytes written to `tmp_path`. Covers: quote-strip (EN/PL/Outlook markers, bare `>` lines), HTML->text (style/script/blockquote/ gmail_quote skipping), newsletter classification, threading extraction, prefix building, body extraction (plain-preferred, HTML fallback, attachment-only), the typed/compat32 parse fallback, stats balance, idempotency (second run inserts nothing new), newsletter chunks -never reaching Ollama, Ollama-offline batch isolation, dimension-mismatch abort, and the -circuit breaker (trips on N consecutive failures, resets on a success, disabled by `0`, -flushes pending threading on abort). +never reaching Ollama, dimension-mismatch abort, and the circuit breaker (trips on N +consecutive give-ups, resets on a success, disabled by `0`, flushes pending threading on +abort). Batch-specific: a poison chunk isolated out of a full batch, a partial failure not +advancing the breaker, and an Ollama hang degrading to the breaker instead of crashing the run. + +The batch client itself is tested in `packages/kb-retrieval/tests/test_embed.py` (22) — retry +with backoff, bisection, dead-backend give-up bounds, and the `TimeoutError`-is-not-a- +`ClientError` regression. The benchmark's pure logic (batch splitting, derived metrics, table +formatting, failure counting) is in `jobs/mail-body-ingest/tests/test_benchmark.py`. + +In a worktree without a preinstalled venv: + +```bash +python3 -m venv /tmp/kbvenv && /tmp/kbvenv/bin/pip install -q pytest pytest-asyncio \ + -e packages/kb-mail/ -e packages/kb-retrieval/ -e jobs/mail-body-ingest/ +/tmp/kbvenv/bin/python -m pytest jobs/mail-body-ingest/tests/ packages/kb-retrieval/tests/ -q +``` diff --git a/kb/services/job-mail-body-ingest.md b/kb/services/job-mail-body-ingest.md index 17152b5..8287dae 100644 --- a/kb/services/job-mail-body-ingest.md +++ b/kb/services/job-mail-body-ingest.md @@ -3,7 +3,7 @@ okf: "0.1" type: service visibility: private status: active -updated: 2026-07-22 +updated: 2026-08-05 links: - ../runbooks/mail-body-ingest-run.md --- @@ -61,9 +61,10 @@ pip install -e jobs/mail-body-ingest/ 8. **Embed + insert**: newsletter chunks are inserted immediately with `excluded_reason='newsletter'`, `embedding=NULL` (no Ollama call, reversible later); everything else is buffered up to `--batch-size` (default 64) and sent through - `kb_retrieval.embed.embed_batch` (`/api/embed`) before inserting. `ON CONFLICT - (envelope_id, chunk_index, model) DO NOTHING` is checked via the command tag, so a - silent no-op counts as `chunks_conflict_skipped`, never `chunks_inserted`. + `kb_retrieval.embed.embed_batch_resilient` (`/api/embed` with `input` as a list) before + inserting. `ON CONFLICT (envelope_id, chunk_index, model) DO NOTHING` is checked via the + command tag, so a silent no-op counts as `chunks_conflict_skipped`, never + `chunks_inserted`. 9. **Threading append** (Decyzja 10): `In-Reply-To`/`References` (angle brackets stripped, matching `envelope.id`'s bare-Message-ID convention) appended as `entities[type=threading]` via the same idempotent `WHERE NOT EXISTS` UPDATE pattern as @@ -93,29 +94,78 @@ counters in the `summary` line, not the exit code. Exit 2 is different — see b |---|---| | 0 | Balanced, zero errors | | 1 | Balanced-but-imperfect (any `parse_errors`/`missing_file`/`read_errors`/`chunks_errors`/`chunks_conflict_skipped`), an unbalanced sum, or an embedding-dimension abort | -| 2 | `--max-embed-failures` consecutive embed batches failed — the embed backend is down; re-run once it is back | +| 2 | `--max-embed-failures` consecutive embed batches gave up — the embed backend is down; re-run once it is back | ## Ollama-offline tolerance and the circuit breaker -A failed `embed_batch()` call is caught per-batch (`aiohttp.ClientError` -> the whole batch, -up to `--batch-size` chunks, counts as `chunks_errors`; the run logs a warning and continues). -Those chunks never enter the idempotency set, so a later re-run retries them automatically — -no separate checkpointing needed. +Three failure modes, three responses — the first two live in +`kb_retrieval.embed.embed_batch_resilient`, the third here: + +1. **Transient blip** — the batch is retried `--embed-retries` times (default 2) with + exponential backoff (`--embed-backoff`, default 1 s, doubled per attempt). Costs seconds, + loses nothing. +2. **One poison chunk** — `/api/embed` is all-or-nothing, so a single pathological text used to + cost the entire 64-chunk batch it happened to land in. When the retries are exhausted but + `/api/tags` says the backend is alive, the batch is bisected until the bad inputs are + isolated; the healthy remainder embeds normally and only the genuinely bad chunks count as + `chunks_errors`. One bad chunk costs ~log2(batch) extra requests instead of 64 embeddings. +3. **Dead backend** — the probe says down, so the batch gives up *without* bisecting (splitting + against a dead backend would burn 2n-1 requests and delay the breaker exactly when it needs + to trip). Tolerating a *flaky* backend is right; surviving a *dead* one is not. The archive is parsed single-threaded ahead of the GPU, so on a full-corpus run (Etap B) a dead Ollama would let the job chew through 200k+ mails at parse speed, mark every chunk `chunks_errors`, and throw away a multi-hour pass. `--max-embed-failures` (default 5, `0` disables) therefore stops the run after -that many *consecutive* failed batches, with exit code 2; a single successful batch resets the -counter. Ollama@SOLARIA's known failure mode is total (container vanishes, network-detached — -4 incidents, plan §1.4/§7), so the breaker trips within seconds of it. On abort, pending -`entities[type=threading]` appends are flushed first: they don't depend on Ollama, they're -idempotent, and re-deriving them would mean re-reading the same 27 GB. +that many *consecutive* give-ups, with exit code 2; any successful batch resets the counter. +Ollama@SOLARIA's known failure mode is total (container vanishes, network-detached — 4 +incidents, plan §1.4/§7), so the breaker trips within seconds of it. On abort, rows already +embedded in the dying batch are committed and pending `entities[type=threading]` appends are +flushed first: they don't depend on Ollama, they're idempotent, and re-deriving them would mean +re-reading the same 27 GB. + +A partial failure against a *live* backend deliberately does **not** advance the breaker: those +chunks were never inserted, so the next run retries them through the ordinary idempotency path, +and letting a handful of bad chunks abort a 50k slice would be strictly worse than skipping +them. `embed_items_failed` and the `embed_requests_total`/`embed_calls` ratio in the `summary` +line are what to read — a ratio of 1.0 means a clean run with no retries or bisection. + +Timeouts are part of this, not an exception to it: aiohttp raises a bare `builtins.TimeoutError` +when `ClientTimeout(total=...)` expires, and that is **not** a subclass of +`aiohttp.ClientError`. An earlier version of this job caught `ClientError` alone, so an Ollama +that accepted the connection and then hung — its actual failure mode — crashed the run outright, +with no breaker and no threading flush. `kb_retrieval.embed.TRANSIENT_EMBED_ERRORS` now names +both; catch that tuple, never `ClientError` on its own. Only a wrong embedding dimension is more severe (`EmbeddingDimensionError`, exit 1) — it aborts immediately, since that would otherwise silently index a vector that doesn't match `document_chunk.embedding VECTOR(1024)`. +## No SOLARIA -> PIHA fallback here (deliberate) + +`kb-query` fails a *query* over to Ollama@PIHA when SOLARIA is down (`app/embed_router.py`). +This job does not, by decision: ~271k active chunks at PIHA's ~790 ms/embed CPU is ~60 h on an +8 GB node already shared with Home Assistant, Paperless and kb-postgres, and CPU embedding does +not batch-scale the way the GPU does. A fallback would quietly turn "abort and resume the slice +once Ollama is fixed" into a two-day run suffocating the infra node. The correct answer to a +dead primary in the backfill path is to stop and resume — which costs nothing, because the job +is idempotent. The two backend paths stay separate on purpose. + +## Batch-size benchmark + +`mail-body-ingest-bench` (same package) sweeps batch sizes over real mail chunks and prints +ms/chunk, chunks/s and a projected wall-clock for the full corpus. Read-only — `SELECT`s and +inference, with no code path that can write — so it is safe to point at live kb-postgres@PIHA. + +```bash +mail-body-ingest-bench --dsn postgresql://kb:@piha:5433/kb \ + --archive-root /home/oskar/kb/mail/archive --sample-envelopes 200 +``` + +It warms the model up before measuring (the first call after an idle period pays the model +load) and measures the *same* chunk set at every size, since ms/chunk depends heavily on text +length. `--max-chunks` caps the set — batch=1 otherwise dominates the wall clock. + ## Idempotency Pre-fetched `(envelope_id, chunk_index)` pairs, scoped server-side to `--model` (`WHERE model @@ -127,7 +177,8 @@ repeating the ambiguity here). ## Definition of Done -Per `CLAUDE.md`: `pytest` passes (55/55) + a `--limit 5` dry-run smoke against live +Per `CLAUDE.md`: `pytest` passes (62/62 for the job, 22/22 for `kb-retrieval`'s embed client) ++ a `--limit 5` dry-run smoke against live `kb-postgres@PIHA` before committing (confirms DSN/query wiring; a missing local archive mirror correctly reports `missing_file` rather than crashing). The Etap A pilot (`--since 2025-07-01 --apply`, plan §7) is a separate, explicitly-confirmed run — not part of this diff --git a/packages/kb-retrieval/src/kb_retrieval/embed.py b/packages/kb-retrieval/src/kb_retrieval/embed.py index 0ae5e06..ff278e0 100644 --- a/packages/kb-retrieval/src/kb_retrieval/embed.py +++ b/packages/kb-retrieval/src/kb_retrieval/embed.py @@ -13,13 +13,43 @@ probing logic lives in exactly one place. vs ~150-200 ms/chunk sequential through `embed_chunk`'s `/api/embeddings` (plan §1.4). Used by `jobs/mail-body-ingest` only -- `embed_chunk` stays the single-text path for kb-query (one query = one text) and the paperless cyclic ingest. + +Two deliberately separate backend paths -- do NOT unify them: + +* **online query** (`kb-query`) routes through `app/embed_router.py`, which fails a query over + to Ollama@PIHA when Ollama@SOLARIA is down. One query = one text = ~790 ms of Pi-5 CPU, an + entirely reasonable price for answering instead of erroring. +* **backfill** (`jobs/mail-body-ingest`) calls this module directly with NO fallback, by + decision. ~271k active chunks at PIHA's 790 ms/embed is ~60 h on an 8 GB node already shared + with Home Assistant, Paperless and kb-postgres, and CPU embedding does not batch-scale the way + the GPU does. A fallback would silently convert "abort and resume the slice once Ollama is + fixed" into a two-day run suffocating the infra node. The correct backfill answer to a dead + primary is to stop (`gave_up` below -> the job's circuit breaker -> exit 2) and resume, which + costs nothing because the job is idempotent. + +`embed_batch_resilient` -- retry + poison-item isolation on top of `embed_batch`, added before +Etap B (the full 225k-envelope pass). It exists because a single transient blip otherwise wrote +off a whole 64-chunk batch, and because `/api/embed` is all-or-nothing: one pathological text +fails the entire batch it happens to land in, indefinitely, on an otherwise healthy backend. """ from __future__ import annotations +import asyncio import time +from dataclasses import dataclass, field import aiohttp +# Transport-level failures that mean "try again", as opposed to a bad payload or a wrong model. +# +# `TimeoutError` is listed EXPLICITLY and must stay: aiohttp raises a bare `builtins.TimeoutError` +# when a `ClientTimeout(total=...)` is exhausted, and that is NOT a subclass of +# `aiohttp.ClientError` (verified on aiohttp 3.14.3 -- only `ServerTimeoutError`, i.e. connect/ +# sock-read timeouts, inherits both). Callers that caught `aiohttp.ClientError` alone therefore +# crashed instead of degrading when Ollama accepted the connection and then hung -- which is +# precisely Ollama@SOLARIA's documented failure mode. Catch this tuple, never `ClientError` alone. +TRANSIENT_EMBED_ERRORS = (aiohttp.ClientError, TimeoutError) + DEFAULT_OLLAMA_URL = "http://localhost:11434" DEFAULT_MODEL = "bge-m3" EXPECTED_DIM = 1024 @@ -59,7 +89,8 @@ async def embed_chunk( async def embed_batch( - session: aiohttp.ClientSession, base_url: str, model: str, texts: list[str] + session: aiohttp.ClientSession, base_url: str, model: str, texts: list[str], + timeout_s: float | None = None, ) -> tuple[list[list[float]], float]: """POST /api/embed on Ollama with a batch of texts. Returns (embeddings, elapsed_seconds). @@ -67,13 +98,24 @@ async def embed_batch( §1.4/§4 decision 5), so parallel requests would only add overhead and unpredictable error ordering in the caller's stats bilans. + `timeout_s` bounds this one request; omitted, the request inherits whatever timeout the + caller's `aiohttp.ClientSession` was constructed with. A per-request bound matters for + batching specifically: batch size and timeout have to move together, and a session-wide + `total` sized for batch 64 is far too generous for batch 1. + Raises `EmbeddingDimensionError` (abort-run) if the response's embedding count doesn't match `len(texts)`, or if any single embedding's dimension isn't `EXPECTED_DIM` -- never silently indexes a vector that doesn't match the `document_chunk.embedding VECTOR(1024)` column. - Same `aiohttp.ClientError` propagation as `embed_chunk` -- no built-in timeout/retry. + Transport failures propagate as `TRANSIENT_EMBED_ERRORS` -- no built-in retry (see + `embed_batch_resilient`). """ + kwargs = {} + if timeout_s is not None: + kwargs["timeout"] = aiohttp.ClientTimeout(total=timeout_s) t0 = time.monotonic() - async with session.post(f"{base_url}/api/embed", json={"model": model, "input": texts}) as resp: + async with session.post( + f"{base_url}/api/embed", json={"model": model, "input": texts}, **kwargs + ) as resp: resp.raise_for_status() data = await resp.json() elapsed = time.monotonic() - t0 @@ -94,6 +136,118 @@ async def embed_batch( return embeddings, elapsed +@dataclass +class BatchOutcome: + """Result of `embed_batch_resilient`. `embeddings[i]` is None exactly for the inputs that + could not be embedded, so the caller can insert the successes and count the rest as errors + without any index bookkeeping of its own. + + `gave_up` is the signal a circuit breaker should count: True means "the backend stopped + answering", False means "the backend is alive and answered for at least part of this batch". + A caller must not treat a partial failure as backend death -- the failed items were never + written, so an idempotent job simply picks them up on the next run. + """ + + embeddings: list[list[float] | None] + elapsed_s: float = 0.0 + requests: int = 0 + retries: int = 0 + failed_indices: list[int] = field(default_factory=list) + gave_up: bool = False + + @property + def ok_count(self) -> int: + return sum(1 for e in self.embeddings if e is not None) + + +async def embed_batch_resilient( + session: aiohttp.ClientSession, + base_url: str, + model: str, + texts: list[str], + *, + retries: int = 2, + backoff_s: float = 1.0, + timeout_s: float | None = None, + health_timeout_s: float = 3.0, + sleep=asyncio.sleep, +) -> BatchOutcome: + """`embed_batch` plus transient-failure retry and poison-item isolation. + + Three failure modes, three distinct responses: + + 1. **Transient blip** -- retried up to `retries` times with exponential backoff + (`backoff_s * 2**attempt`). Costs a few seconds, loses nothing. + 2. **Backend down** -- after the retries are exhausted, `/api/tags` is probed. If the probe + says down, the whole span is marked failed and `gave_up=True` is returned immediately. + No bisection: splitting a batch against a dead backend would burn 2n-1 requests and + delay the caller's circuit breaker exactly when it most needs to trip. + 3. **Poison item** -- probe says up, so the backend is answering and one of the inputs is + the problem. Bisect and recurse to find it. Healthy halves are embedded normally and + only the genuinely bad inputs land in `failed_indices`; one bad chunk costs its batch + ~log2(n) extra requests instead of all n embeddings. + + Bisected sub-spans do NOT retry (`retries=0`). The top-level attempt already established + that this is not a transient blip, and re-probing health at every failed span means a + backend that dies mid-bisection still short-circuits to `gave_up` rather than grinding on. + Worst case -- every input poisoned against a live backend -- is 2n-1 requests, which buys + the exact identity of all n bad chunks and is a pathology worth paying once to see. + + `EmbeddingDimensionError` and `ValueError` are never retried or swallowed: a wrong vector + dimension or a malformed response is a loud config/version fault, not something to degrade + around. They propagate to the caller. + """ + outcome = BatchOutcome(embeddings=[None] * len(texts)) + if not texts: + return outcome + + async def attempt(span_texts: list[str], span_retries: int) -> list[list[float]] | None: + """One span, retried `span_retries` times. None = still failing after the retries.""" + for n in range(span_retries + 1): + try: + outcome.requests += 1 + embeddings, elapsed = await embed_batch( + session, base_url, model, span_texts, timeout_s=timeout_s + ) + outcome.elapsed_s += elapsed + return embeddings + except TRANSIENT_EMBED_ERRORS: + if n == span_retries: + return None + outcome.retries += 1 + await sleep(backoff_s * (2**n)) + return None + + async def resolve(indices: list[int], span_retries: int) -> None: + if outcome.gave_up: + return + span_texts = [texts[i] for i in indices] + + embeddings = await attempt(span_texts, span_retries) + if embeddings is not None: + for i, embedding in zip(indices, embeddings): + outcome.embeddings[i] = embedding + return + + if not await check_ollama_health(session, base_url, health_timeout_s): + outcome.gave_up = True + return + + if len(indices) == 1: + return + + mid = len(indices) // 2 + await resolve(indices[:mid], 0) + await resolve(indices[mid:], 0) + + await resolve(list(range(len(texts))), retries) + # Derived, not accumulated: when `gave_up` trips mid-bisection the abandoned sub-spans never + # reach a leaf, so any per-span bookkeeping would silently under-report them. The embeddings + # list is the single source of truth -- unset means failed, on every path. + outcome.failed_indices = [i for i, e in enumerate(outcome.embeddings) if e is None] + return outcome + + async def check_ollama_health( session: aiohttp.ClientSession, base_url: str, timeout_s: float ) -> bool: diff --git a/packages/kb-retrieval/tests/test_embed.py b/packages/kb-retrieval/tests/test_embed.py index adf7cf7..e1bef70 100644 --- a/packages/kb-retrieval/tests/test_embed.py +++ b/packages/kb-retrieval/tests/test_embed.py @@ -1,13 +1,16 @@ """Unit tests for the Ollama embedding client -- no real HTTP, no real Ollama.""" from __future__ import annotations +import aiohttp import pytest from kb_retrieval.embed import ( + TRANSIENT_EMBED_ERRORS, EmbeddingDimensionError, _vector_literal, check_ollama_health, embed_batch, + embed_batch_resilient, embed_chunk, ) @@ -86,8 +89,8 @@ class _FakeBatchOllamaSession: self._bad_index = bad_index self.requests: list[dict] = [] - def post(self, url, json): - self.requests.append({"url": url, "json": json}) + def post(self, url, json, timeout=None): + self.requests.append({"url": url, "json": json, "timeout": timeout}) texts = json["input"] n = self._n_override if self._n_override is not None else len(texts) embeddings = [[0.01] * self._dim for _ in range(n)] @@ -105,9 +108,15 @@ class TestEmbedBatch: assert all(len(e) == 1024 for e in embeddings) assert elapsed >= 0 assert session.requests == [ - {"url": "http://fake-ollama/api/embed", "json": {"model": "bge-m3", "input": texts}} + {"url": "http://fake-ollama/api/embed", + "json": {"model": "bge-m3", "input": texts}, "timeout": None} ] + async def test_timeout_s_bounds_the_single_request(self): + session = _FakeBatchOllamaSession(dim=1024) + await embed_batch(session, "http://fake-ollama", "bge-m3", ["a"], timeout_s=5.0) + assert session.requests[0]["timeout"].total == 5.0 + async def test_length_mismatch_raises_dimension_error(self): session = _FakeBatchOllamaSession(dim=1024, n_override=2) with pytest.raises(EmbeddingDimensionError): @@ -167,3 +176,186 @@ class TestCheckOllamaHealth: async def test_down_on_timeout(self): session = _FakeHealthSession(raises=TimeoutError()) assert await check_ollama_health(session, "http://fake-ollama", 0.5) is False + + +class _FailResponse(_FakeEmbedResponse): + """A transport failure, i.e. something in TRANSIENT_EMBED_ERRORS -- not the plain + RuntimeError `_FakeEmbedResponse` raises, which would (correctly) never be retried.""" + + def raise_for_status(self): + raise aiohttp.ClientConnectionError("simulated transport failure") + + +class _ResilientSession: + """Drives every branch of `embed_batch_resilient`. + + * `poison` -- texts that fail whatever batch they land in (the all-or-nothing property of + /api/embed that bisection exists to work around) + * `fail_first` -- the first N posts fail regardless (a transient blip) + * `post_exc` -- raised instead of returning a response (used for the TimeoutError regression) + * `health_sequence` -- consumed one entry per /api/tags probe, then falls back to `health` + """ + + def __init__(self, *, dim=1024, poison=(), fail_first=0, health=True, + health_sequence=None, post_exc=None): + self._dim = dim + self._poison = set(poison) + self._fail_first = fail_first + self._health = health + self._health_sequence = list(health_sequence or []) + self._post_exc = post_exc + self.posts: list[list[str]] = [] + self.health_checks = 0 + + def post(self, url, json, timeout=None): + texts = json["input"] + self.posts.append(list(texts)) + if self._post_exc is not None: + raise self._post_exc + if self._fail_first > 0: + self._fail_first -= 1 + return _FailResponse({}) + if self._poison & set(texts): + return _FailResponse({}) + return _FakeBatchResponse({"embeddings": [[0.01] * self._dim for _ in texts]}) + + def get(self, url, timeout=None): + self.health_checks += 1 + up = self._health_sequence.pop(0) if self._health_sequence else self._health + return _FakeTagsResponse(200 if up else 500) + + +def _recorder(): + sleeps: list[float] = [] + + async def fake_sleep(seconds): + sleeps.append(seconds) + + return sleeps, fake_sleep + + +class TestEmbedBatchResilient: + async def test_clean_run_costs_one_request_and_no_health_probe(self): + session = _ResilientSession() + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", ["a", "b", "c"] + ) + assert outcome.ok_count == 3 + assert outcome.failed_indices == [] + assert outcome.gave_up is False + assert outcome.requests == 1 + assert outcome.retries == 0 + assert session.health_checks == 0 + + async def test_transient_failure_is_retried_with_exponential_backoff(self): + session = _ResilientSession(fail_first=2) + sleeps, fake_sleep = _recorder() + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", ["a", "b"], + retries=2, backoff_s=1.0, sleep=fake_sleep, + ) + assert outcome.ok_count == 2 + assert outcome.requests == 3 + assert outcome.retries == 2 + assert sleeps == [1.0, 2.0] + assert session.health_checks == 0 # never needed -- the retry recovered it + + async def test_timeout_is_transient_and_never_escapes(self): + """Regression: aiohttp raises a bare builtins.TimeoutError when ClientTimeout(total=...) + expires, and that is NOT an aiohttp.ClientError. Callers catching ClientError alone + crashed instead of degrading -- exactly Ollama@SOLARIA's hang-not-refuse failure mode.""" + assert not isinstance(TimeoutError(), aiohttp.ClientError) # the trap this guards + assert isinstance(TimeoutError(), TRANSIENT_EMBED_ERRORS) + + session = _ResilientSession(post_exc=TimeoutError(), health=False) + _sleeps, fake_sleep = _recorder() + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", ["a", "b"], retries=1, sleep=fake_sleep, + ) + assert outcome.gave_up is True + assert outcome.failed_indices == [0, 1] + + async def test_dead_backend_gives_up_without_bisecting(self): + texts = [f"t{i}" for i in range(8)] + session = _ResilientSession(poison=texts, health=False) + _sleeps, fake_sleep = _recorder() + + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", texts, retries=1, sleep=fake_sleep, + ) + + assert outcome.gave_up is True + assert outcome.ok_count == 0 + assert outcome.failed_indices == list(range(8)) + # 1 attempt + 1 retry, then the health probe stops it. Bisecting a dead backend would + # cost 2n-1 requests and delay the caller's circuit breaker. + assert len(session.posts) == 2 + assert session.health_checks == 1 + + async def test_poison_item_is_isolated_and_the_rest_still_embed(self): + texts = [f"t{i}" for i in range(8)] + session = _ResilientSession(poison={"t3"}, health=True) + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", texts, retries=0, + ) + + assert outcome.gave_up is False + assert outcome.failed_indices == [3] + assert outcome.ok_count == 7 + assert outcome.embeddings[3] is None + assert all(outcome.embeddings[i] is not None for i in range(8) if i != 3) + assert len(session.posts) < 2 * len(texts) # bisection, not one request per item + + async def test_bisected_spans_do_not_retry(self): + """The top-level attempt already proved this isn't a blip; re-retrying every sub-span + would multiply a poison chunk's cost by (retries+1) at every level of the bisection.""" + texts = ["a", "b", "c", "d"] + session = _ResilientSession(poison={"c"}, health=True) + sleeps, fake_sleep = _recorder() + + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", texts, + retries=2, backoff_s=1.0, sleep=fake_sleep, + ) + + assert outcome.failed_indices == [2] + assert sleeps == [1.0, 2.0] # top level only -- no backoff inside the bisection + assert outcome.retries == 2 + + async def test_backend_dying_mid_bisection_fails_every_unresolved_item(self): + """The abandoned right-hand span must still be reported failed. Accumulating + failed_indices per span silently under-reported it -- it never reaches a leaf.""" + texts = ["a", "b", "c", "d"] + session = _ResilientSession(poison=texts, health_sequence=[True, False]) + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", texts, retries=0, + ) + + assert outcome.gave_up is True + assert outcome.ok_count == 0 + assert outcome.failed_indices == [0, 1, 2, 3] + assert all(e is None for e in outcome.embeddings) + + async def test_dimension_error_is_never_retried_or_swallowed(self): + session = _FakeBatchOllamaSession(dim=1024, n_override=2) + with pytest.raises(EmbeddingDimensionError): + await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", ["a", "b", "c"], retries=3, + ) + assert len(session.requests) == 1 + + async def test_empty_input_makes_no_requests(self): + session = _ResilientSession() + outcome = await embed_batch_resilient(session, "http://fake-ollama", "bge-m3", []) + assert outcome.embeddings == [] + assert outcome.requests == 0 + assert session.posts == [] + + async def test_elapsed_and_requests_are_accumulated_across_retries(self): + session = _ResilientSession(fail_first=1) + _sleeps, fake_sleep = _recorder() + outcome = await embed_batch_resilient( + session, "http://fake-ollama", "bge-m3", ["a"], retries=1, sleep=fake_sleep, + ) + assert outcome.requests == 2 + assert outcome.elapsed_s >= 0