Compare commits
2 commits
71eaab0025
...
75116ad0a5
| Author | SHA1 | Date | |
|---|---|---|---|
|
75116ad0a5 | ||
|
|
02a00797f4 |
|
|
@ -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"]
|
||||
|
|
|
|||
241
jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py
Normal file
241
jobs/mail-body-ingest/src/mail_body_ingest/benchmark.py
Normal file
|
|
@ -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:<pw>@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()
|
||||
|
|
@ -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"] == (
|
||||
|
|
|
|||
155
jobs/mail-body-ingest/tests/test_benchmark.py
Normal file
155
jobs/mail-body-ingest/tests/test_benchmark.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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:<pw>@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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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:<pw>@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
|
||||
|
|
|
|||
|
|
@ -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,29 @@ 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`).
|
||||
|
||||
No backend failover: this is the backfill path, and it must NOT fall back to Ollama@PIHA the
|
||||
way `kb-query`'s `embed_router` does for online queries. ~271k chunks at PIHA's ~790 ms/embed
|
||||
CPU is ~60 h on an 8 GB node shared with Home Assistant, Paperless and kb-postgres. See the
|
||||
module docstring for the full argument -- the two paths are separate by decision.
|
||||
"""
|
||||
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 +141,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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue