diff --git a/Makefile b/Makefile index ea6fa5e..d33f305 100644 --- a/Makefile +++ b/Makefile @@ -141,6 +141,14 @@ stats: bootstrap ## counts: documents, chunks, edges, audit chain test: bootstrap ## run pytest suite $(VENV)/bin/pytest -q +# Reproducible micro-benchmark over a fixed slice of cur. Lets you compare +# ETL throughput across configs and catches regressions on optimization +# work. Override BENCH_DOCS=N (default 5000). +BENCH_DOCS ?= 5000 +BENCH_DIR := /tmp/aborist-bench +bench: bootstrap fetch-cur ## benchmark serial vs parallel-shared vs attached at $(BENCH_DOCS) docs + @bash bench/run.sh $(BENCH_DOCS) + clean: ## remove venv + caches (keeps fetched data and db) rm -rf $(VENV) .pytest_cache **/__pycache__ aborist.egg-info find . -type d -name __pycache__ -prune -exec rm -rf {} + diff --git a/aborist/cli.py b/aborist/cli.py index 7e551fc..405bca0 100644 --- a/aborist/cli.py +++ b/aborist/cli.py @@ -9,6 +9,7 @@ from pathlib import Path from aborist import __version__ from aborist.ingest import ingest_source, verify_random_sample +from aborist.progress import Progress from aborist.search import FTS5Backend from aborist.sources import WikipediaCurDump from aborist.store import DEFAULT_DB_PATH, connect, connect_query, stats @@ -70,6 +71,17 @@ def _cmd_ingest(args: argparse.Namespace) -> int: digits = max(3, len(str(total - 1))) target_db = shards_dir / f"{rank:0{digits}d}.db" + progress: Progress | None = None + if not args.quiet: + prefix = "" + if args.shard: + prefix = f"[shard {args.shard}] " + progress = Progress( + interval=args.progress_interval, + total_estimate=args.total_estimate, + prefix=prefix, + ) + conn = connect(target_db) try: result = ingest_source( @@ -79,6 +91,7 @@ def _cmd_ingest(args: argparse.Namespace) -> int: limit=args.limit, batch_size=args.batch_size, resume=args.resume, + progress=progress, ) finally: conn.close() @@ -593,6 +606,28 @@ def build_parser() -> argparse.ArgumentParser: "and restart at any time" ), ) + ingest.add_argument( + "--quiet", + action="store_true", + help="suppress periodic stderr progress output", + ) + ingest.add_argument( + "--progress-interval", + dest="progress_interval", + type=float, + default=2.0, + help="seconds between stderr progress lines (default 2.0)", + ) + ingest.add_argument( + "--total-estimate", + dest="total_estimate", + type=int, + default=None, + help=( + "estimated total docs the source will yield. enables percent " + "+ ETA in progress output" + ), + ) ingest.set_defaults(func=_cmd_ingest) search = sub.add_parser("search", help="keyword search (VISUAL audit mode)") diff --git a/aborist/ingest.py b/aborist/ingest.py index ff39d95..67216ce 100644 --- a/aborist/ingest.py +++ b/aborist/ingest.py @@ -23,6 +23,7 @@ from aborist import ( ) from aborist.document import Document, canonicalize, get_chunker from aborist.merkle import MerkleTree, hash_leaf +from aborist.progress import Progress from aborist.source import Source from aborist.store import ( chain_audit_events, @@ -60,6 +61,7 @@ def ingest_source( limit: int | None = None, batch_size: int = DEFAULT_BATCH_SIZE, resume: bool = False, + progress: Progress | None = None, ) -> IngestStats: """Ingest every document the source yields. Returns counts. @@ -67,6 +69,9 @@ def ingest_source( table and asks the source to fast-forward past it. After each successful batch flush, the high-water mark is updated in meta. A killed process can rsync forward by re-running with --resume. + + `progress` (optional) gets a `tick(seen, inserted=...)` call after each + batch flush. Pass an `aborist.progress.Progress` for live stderr output. """ chunker = get_chunker(chunker_name) stats = IngestStats() @@ -89,11 +94,11 @@ def ingest_source( stats.inserted += inserted stats.skipped_duplicate += skipped batch.clear() - # Persist high-water after each successful batch so a kill anywhere - # past this point lets a future --resume skip ahead cheaply. if hasattr(source, "last_id") and source.last_id: with transaction(conn): set_meta(conn, meta_key, str(source.last_id)) + if progress is not None: + progress.tick(stats.seen, inserted=stats.inserted) for doc in source.iter_documents(): stats.seen += 1 @@ -108,6 +113,9 @@ def ingest_source( flush() flush() + if progress is not None: + progress.done(stats.seen, inserted=stats.inserted) + stats.chunks_total = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0] stats.edges_total = conn.execute("SELECT COUNT(*) FROM edges").fetchone()[0] return stats diff --git a/aborist/progress.py b/aborist/progress.py new file mode 100644 index 0000000..70f918d --- /dev/null +++ b/aborist/progress.py @@ -0,0 +1,86 @@ +"""Periodic stderr progress reporting for long-running ingests/distillation. + +Stdlib only. Prints on a rate-limited schedule (default every 2 seconds) so +ingests of 100k+ docs give live feedback without spamming. Optional total +estimate produces a percent + ETA. + +Multiple parallel shard processes will interleave their lines; each can be +given a `prefix` to disambiguate. +""" + +from __future__ import annotations + +import sys +import time +from typing import IO + + +class Progress: + """Rate-limited progress reporter to stderr (or any stream).""" + + def __init__( + self, + *, + interval: float = 2.0, + total_estimate: int | None = None, + prefix: str = "", + stream: IO[str] = sys.stderr, + ): + self.start = time.monotonic() + self.last_print = self.start + self.last_seen = 0 + self.interval = interval + self.total_estimate = total_estimate + self.prefix = prefix + self.stream = stream + + def tick( + self, + seen: int, + *, + inserted: int | None = None, + force: bool = False, + ) -> None: + """Maybe print a status line. Call after each batch flush.""" + now = time.monotonic() + if not force and (now - self.last_print) < self.interval: + return + elapsed = now - self.start + delta_seen = max(0, seen - self.last_seen) + delta_t = max(now - self.last_print, 1e-6) + rate_now = delta_seen / delta_t + rate_avg = seen / elapsed if elapsed > 0 else 0.0 + + bits = [ + f"{self.prefix}[{_format_elapsed(elapsed)}]", + f"{seen:>10,} seen", + ] + if inserted is not None: + bits.append(f"{inserted:>9,} new") + bits.append(f"{rate_now:>5.0f} now") + bits.append(f"({rate_avg:>4.0f} avg) docs/s") + if self.total_estimate and rate_now > 0: + pct = 100.0 * min(1.0, seen / self.total_estimate) + remaining = max(0, self.total_estimate - seen) + eta = remaining / rate_now + bits.append(f"{pct:>5.1f}%") + bits.append(f"ETA {_format_elapsed(eta)}") + + print(" " + " | ".join(bits), file=self.stream, flush=True) + self.last_print = now + self.last_seen = seen + + def done(self, seen: int, *, inserted: int | None = None) -> None: + """Print a final status line regardless of interval.""" + self.tick(seen, inserted=inserted, force=True) + + +def _format_elapsed(secs: float) -> str: + secs = int(secs) + h, rem = divmod(secs, 3600) + m, s = divmod(rem, 60) + if h: + return f"{h}h{m:02d}m" + if m: + return f"{m:>2}m{s:02d}s" + return f" {s:>2}s" diff --git a/bench/run.sh b/bench/run.sh new file mode 100755 index 0000000..8fba13c --- /dev/null +++ b/bench/run.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# Reproducible ETL throughput benchmark. +# +# Three configurations against a fixed BENCH_DOCS slice of the Wikipedia +# 2003-05-16 cur dump: +# +# serial single process, single SQLite (Phase 1a) +# parallel-shared N shards, one shared SQLite (Phase 1b — WAL serialized) +# attached N shards, per-shard SQLite (Phase 2 — true parallel writes) +# +# Output: per-config wall, throughput (docs/s), and a one-line summary table. +set -uo pipefail + +BENCH_DOCS=${1:-5000} +BENCH_DIR=${BENCH_DIR:-/tmp/aborist-bench} +SHARDS=${SHARDS:-4} +DUMP=${DUMP:-data/20030516_cur_tablesql.bz2} +ABORIST=${ABORIST:-.venv/bin/aborist} + +mkdir -p "$BENCH_DIR" + +if [[ ! -f "$DUMP" ]]; then + echo "missing dump $DUMP — run 'make fetch-cur' first" >&2 + exit 2 +fi + +run_silent() { + # $1 = label (writes wall time in seconds to stdout) + local start end + start=$(date +%s.%N) + "$@" >/dev/null 2>&1 + end=$(date +%s.%N) + echo "$end - $start" | bc -l +} + +# Each row will be: label,wall_s,docs/s +declare -a results + +bench_serial() { + local db=$BENCH_DIR/serial.db + rm -f "$db" "$db-"* + local wall + wall=$(run_silent "$ABORIST" --db "$db" ingest --quiet \ + --source wikipedia_cur --path "$DUMP" --limit "$BENCH_DOCS") + local docs + docs=$(sqlite3 "$db" "SELECT COUNT(*) FROM documents") + printf "serial,%s,%s\n" "$wall" "$docs" +} + +bench_parallel_shared() { + local db=$BENCH_DIR/shared.db + rm -f "$db" "$db-"* + local per=$((BENCH_DOCS / SHARDS)) + local start end wall + start=$(date +%s.%N) + for i in $(seq 0 $((SHARDS - 1))); do + "$ABORIST" --db "$db" ingest --quiet \ + --source wikipedia_cur --path "$DUMP" \ + --shard "$i/$SHARDS" --limit "$per" >/dev/null 2>&1 & + done + wait + end=$(date +%s.%N) + wall=$(echo "$end - $start" | bc -l) + local docs + docs=$(sqlite3 "$db" "SELECT COUNT(*) FROM documents") + printf "parallel-shared,%s,%s\n" "$wall" "$docs" +} + +bench_attached() { + local dir=$BENCH_DIR/attached + rm -rf "$dir" + mkdir -p "$dir" + local per=$((BENCH_DOCS / SHARDS)) + local start end wall + start=$(date +%s.%N) + for i in $(seq 0 $((SHARDS - 1))); do + "$ABORIST" ingest --quiet \ + --source wikipedia_cur --path "$DUMP" \ + --shards-dir "$dir" --shard "$i/$SHARDS" --limit "$per" >/dev/null 2>&1 & + done + wait + end=$(date +%s.%N) + wall=$(echo "$end - $start" | bc -l) + local docs + docs=$("$ABORIST" --shards-dir "$dir" stats 2>/dev/null \ + | python3 -c "import json,sys; print(json.load(sys.stdin)['documents_total'])") + printf "attached,%s,%s\n" "$wall" "$docs" +} + +echo "=== aborist ETL benchmark — $BENCH_DOCS docs target, $SHARDS shards ===" +echo + +results+=("$(bench_serial)") +results+=("$(bench_parallel_shared)") +results+=("$(bench_attached)") + +# Pretty table. +printf " %-18s %10s %10s %12s\n" "config" "wall_s" "docs" "docs/s" +printf " %-18s %10s %10s %12s\n" "------" "------" "----" "------" +for row in "${results[@]}"; do + IFS=',' read -r label wall docs <<<"$row" + rate=$(echo "scale=1; $docs / $wall" | bc -l) + printf " %-18s %10.2f %10s %12s\n" "$label" "$wall" "$docs" "$rate" +done +echo +echo " (raw CSV: $BENCH_DIR/results.csv)" +{ printf "config,wall_s,docs\n"; printf "%s\n" "${results[@]}"; } > "$BENCH_DIR/results.csv"