arborist/bench/shard_count_sweep.py

#!/usr/bin/env python3
"""Shard-count sweep — measure ingest + query cost as a function of M.

Variables:
    M ∈ {1, 2, 4, 8}     shard count for that run
    N = M                ingest workers (one per shard, no WAL contention)

Fixed:
    corpus = N_DOCS synthetic deterministic-content documents
    queries = N_QUERIES fixed FTS5 queries (head + tail tokens)

Output:
    - CSV at bench/results/shard-count-sweep-<ts>.csv
    - Console summary table

The script wipes ~/.arborist/bench-shards/ between runs so it never
touches the live corpus at ~/.arborist/shards/.

This benchmark exists to answer one question for #000065:
**at what M does ingest throughput stop improving, and how does query
latency change?** Answer drives the canonical-M choice.
"""

from __future__ import annotations

import csv
import hashlib
import json
import os
import random
import shutil
import statistics
import subprocess
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

# Make arborist importable from the repo.
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from arborist.document import Document, canonicalize  # noqa: E402
from arborist.ingest import ingest_source  # noqa: E402
from arborist.merkle import MerkleTree, hash_leaf  # noqa: E402
from arborist.source import Source  # noqa: E402
from arborist.sources.wikipedia import WikipediaCurDump  # noqa: E402
from arborist.store import connect, connect_query, discover_shards  # noqa: E402

BENCH_DIR = Path.home() / ".arborist" / "bench-shards"
RESULTS_DIR = ROOT / "bench" / "results"
WIKI_CUR_PATH = ROOT / "data" / "20030516_cur_tablesql.bz2"

# Compact for time budget. Each cell of the M-sweep re-ingests N_DOCS,
# and Wikipedia ingest does real wikitext canonicalization + edge
# extraction so per-doc cost is realistic.
N_DOCS = 2000
N_QUERIES = 50
SHARD_COUNTS = [1, 2, 4, 8]
USE_WIKIPEDIA = True  # False = synthetic fallback (kept for fast iteration)


# ---------------------------------------------------------------------------
# Synthetic corpus generation
# ---------------------------------------------------------------------------


WORDS = [
    "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel",
    "india", "juliet", "kilo", "lima", "mike", "november", "oscar", "papa",
    "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey",
    "xray", "yankee", "zulu",
    "computer", "wikipedia", "physics", "history", "music", "biology",
    "mathematics", "philosophy", "economics", "literature", "chemistry",
    "geology", "astronomy", "psychology", "sociology",
]


def load_wikipedia_corpus(n_docs: int) -> list[Document]:
    """Read the first n_docs articles from the 2003 Wikipedia cur dump.

    Real wikitext, real canonicalization cost — this is the bench cell
    that matters for the M-decision because synthetic random-words docs
    underestimate ingest per-doc work.
    """
    if not WIKI_CUR_PATH.exists():
        raise FileNotFoundError(
            f"Wikipedia dump not found at {WIKI_CUR_PATH}. "
            f"Run `make fetch-cur` first."
        )
    src = WikipediaCurDump(WIKI_CUR_PATH)
    docs = []
    for doc in src.iter_documents():
        docs.append(doc)
        if len(docs) >= n_docs:
            break
    return docs


def gen_corpus(n_docs: int, seed: int = 42) -> list[Document]:
    """Generate N deterministic synthetic docs. Same seed → same content →
    same document_root across runs. Routes deterministically by hash."""
    rng = random.Random(seed)
    docs = []
    for i in range(n_docs):
        # Each doc ~3 KB of text drawn from word pool. Common words appear
        # more often (so FTS queries on them hit many docs), some rare
        # tokens (doc_NNN tags) appear in exactly one doc.
        sentences = []
        for _ in range(60):
            length = rng.randint(8, 20)
            words = rng.choices(WORDS, k=length)
            sentences.append(" ".join(words).capitalize() + ".")
        # Inject a unique rare token per doc so we can do "find doc X" queries.
        sentences.append(f"This is the unique marker bench_doc_{i:04d}.")
        content = " ".join(sentences)
        docs.append(Document(
            uri=f"bench://doc_{i:04d}",
            content=content,
            source_type="html",   # so license_class is "unknown" but ingest works
            title=f"Bench Doc {i:04d}",
        ))
    return docs


def doc_to_shard(doc: Document, M: int) -> int:
    """Content-hash routing: deterministic shard index from canonicalized
    content. Matches the #000065 routing function — first 8 hex chars of
    document_root (= leaf-Merkle root over canonical chunks)."""
    # Approximate document_root using a stand-in: hash of canonical content
    # bytes. For the bench, what matters is determinism + uniform spread,
    # not exact match to ingest's hash.
    h = hashlib.sha256(canonicalize(doc.content).encode("utf-8")).hexdigest()
    return int(h[:8], 16) % M


class ListSource(Source):
    source_type = "html"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self):
        yield from self.docs


# ---------------------------------------------------------------------------
# Ingest worker — runs in a subprocess so SQLite writer locks are real
# ---------------------------------------------------------------------------


def _ingest_worker(args: tuple[str, list[dict]]) -> dict:
    """Subprocess worker: ingest the assigned docs into one shard.
    args = (db_path, list of {uri, content, title} dicts).
    Returns timing + RSS info.
    """
    import resource
    # Bench infrastructure recycles shard paths between runs. Clear the
    # migration cache so each iteration re-runs SCHEMA_SQL on the fresh
    # (just-deleted) DB. Without this, the M=2+ iterations see a path
    # in _MIGRATED_SHARDS and skip schema creation → "no such table".
    from arborist import store as _store
    _store._MIGRATED_SHARDS.clear()
    db_path, doc_dicts = args
    docs = [
        Document(
            uri=d["uri"], content=d["content"],
            source_type="html", title=d["title"],
        )
        for d in doc_dicts
    ]
    src = ListSource(docs)
    t0 = time.perf_counter()
    conn = connect(db_path)
    try:
        result = ingest_source(conn, src)
    finally:
        conn.close()
    wall = time.perf_counter() - t0
    rss_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return {
        "db": db_path,
        "wall_s": wall,
        "docs_ingested": result.inserted,
        "chunks_ingested": result.chunks_total,
        "peak_rss_mb": rss_kb / 1024,
    }


# ---------------------------------------------------------------------------
# Bench loop
# ---------------------------------------------------------------------------


def reset_bench_dir() -> None:
    if BENCH_DIR.exists():
        shutil.rmtree(BENCH_DIR)
    BENCH_DIR.mkdir(parents=True)
    # Also clear the parent process's migration cache; the in-process
    # M=1 path uses it.
    from arborist import store as _store
    _store._MIGRATED_SHARDS.clear()


def bench_ingest(M: int, docs: list[Document]) -> dict:
    """Run ingest with M shards + M parallel workers. Returns aggregate stats."""
    reset_bench_dir()
    # Hash-partition docs into M groups.
    groups: list[list[Document]] = [[] for _ in range(M)]
    for doc in docs:
        groups[doc_to_shard(doc, M)].append(doc)
    sizes = [len(g) for g in groups]
    # Each worker gets one shard.
    db_paths = [str(BENCH_DIR / f"{i:03d}.db") for i in range(M)]
    args_list = [
        (db_paths[i], [
            {"uri": d.uri, "content": d.content, "title": d.title}
            for d in groups[i]
        ])
        for i in range(M)
    ]
    t0 = time.perf_counter()
    worker_results = []
    if M == 1:
        # No need for process pool when there's one shard.
        worker_results.append(_ingest_worker(args_list[0]))
    else:
        with ProcessPoolExecutor(max_workers=M) as pool:
            futures = [pool.submit(_ingest_worker, args) for args in args_list]
            for fut in as_completed(futures):
                worker_results.append(fut.result())
    wall = time.perf_counter() - t0
    total_chunks = sum(r["chunks_ingested"] for r in worker_results)
    return {
        "M": M,
        "phase": "ingest",
        "wall_s": wall,
        "per_shard_size_min": min(sizes),
        "per_shard_size_max": max(sizes),
        "total_chunks": total_chunks,
        "per_worker_wall_s_max": max(r["wall_s"] for r in worker_results),
        "per_worker_wall_s_min": min(r["wall_s"] for r in worker_results),
        "per_worker_peak_rss_mb_max": max(r["peak_rss_mb"] for r in worker_results),
    }


def gen_queries(rng_seed: int = 7) -> list[str]:
    """Mix of head-term (common, hits many docs) and tail-term (one doc) queries."""
    rng = random.Random(rng_seed)
    queries = []
    # Half head queries: random common words.
    for _ in range(N_QUERIES // 2):
        n_words = rng.randint(1, 3)
        queries.append(" ".join(rng.choices(WORDS, k=n_words)))
    # Half tail queries: a specific bench_doc_NNNN marker.
    sample_docs = rng.sample(range(N_DOCS), N_QUERIES - N_QUERIES // 2)
    for i in sample_docs:
        queries.append(f"bench_doc_{i:04d}")
    return queries


def bench_query(M: int, queries: list[str]) -> dict:
    """Run query suite against the M-shard layout."""
    # Use connect_query which ATTACHes all the shards.
    snap_shards = discover_shards(BENCH_DIR)
    if not snap_shards:
        return {"M": M, "phase": "query", "wall_s": 0, "p50_ms": 0, "p99_ms": 0,
                "queries": 0, "error": "no_shards_discovered"}
    t_attach_0 = time.perf_counter()
    conn = connect_query(":memory:", shards_dir=BENCH_DIR)
    attach_ms = (time.perf_counter() - t_attach_0) * 1000
    try:
        latencies_ms = []
        t_total_0 = time.perf_counter()
        for q in queries:
            t0 = time.perf_counter()
            # FTS5 query across the union view.
            rows = conn.execute(
                "SELECT COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH ?",
                (q,),
            ).fetchall()
            latencies_ms.append((time.perf_counter() - t0) * 1000)
        total_wall = time.perf_counter() - t_total_0
    finally:
        conn.close()
    latencies_ms.sort()
    return {
        "M": M,
        "phase": "query",
        "wall_s": total_wall,
        "attach_ms": attach_ms,
        "p50_ms": latencies_ms[len(latencies_ms) // 2],
        "p99_ms": latencies_ms[int(len(latencies_ms) * 0.99)],
        "mean_ms": statistics.mean(latencies_ms),
        "queries": len(queries),
    }


def main():
    print(f"== shard-count sweep ==")
    print(f"  N_DOCS={N_DOCS}, N_QUERIES={N_QUERIES}, M ∈ {SHARD_COUNTS}")
    print(f"  bench dir: {BENCH_DIR}")
    print()

    if USE_WIKIPEDIA:
        print(f"  corpus: real Wikipedia ({WIKI_CUR_PATH.name}, first {N_DOCS} docs)")
        t_load = time.perf_counter()
        docs = load_wikipedia_corpus(N_DOCS)
        print(f"  loaded {len(docs)} docs in {time.perf_counter() - t_load:.1f}s")
    else:
        print(f"  corpus: synthetic ({N_DOCS} docs)")
        docs = gen_corpus(N_DOCS)
    queries = gen_queries()
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    ts = time.strftime("%Y-%m-%dT%H-%M-%SZ", time.gmtime())
    csv_path = RESULTS_DIR / f"shard-count-sweep-{ts}.csv"

    all_rows = []
    summary = []
    for M in SHARD_COUNTS:
        print(f"-- M={M} --")
        ingest = bench_ingest(M, docs)
        print(f"  ingest:  {ingest['wall_s']:.2f}s  "
              f"chunks={ingest['total_chunks']}  "
              f"worker_max={ingest['per_worker_wall_s_max']:.2f}s  "
              f"peak_rss_mb={ingest['per_worker_peak_rss_mb_max']:.1f}")
        all_rows.append(ingest)
        query = bench_query(M, queries)
        print(f"  query:   {query['wall_s']:.2f}s for {query['queries']} qs  "
              f"p50={query['p50_ms']:.2f}ms  p99={query['p99_ms']:.2f}ms  "
              f"attach={query['attach_ms']:.1f}ms")
        all_rows.append(query)
        summary.append({
            "M": M,
            "ingest_wall_s": ingest["wall_s"],
            "ingest_chunks_per_s": ingest["total_chunks"] / ingest["wall_s"],
            "query_p50_ms": query["p50_ms"],
            "query_p99_ms": query["p99_ms"],
            "query_attach_ms": query["attach_ms"],
        })

    # Write CSV
    with open(csv_path, "w") as f:
        if all_rows:
            keys = set()
            for r in all_rows:
                keys.update(r.keys())
            w = csv.DictWriter(f, fieldnames=sorted(keys))
            w.writeheader()
            for r in all_rows:
                w.writerow(r)
    print()
    print(f"== summary ==")
    print(f"{'M':>2}  {'ingest_wall_s':>14}  {'chunks/s':>10}  "
          f"{'q_p50_ms':>10}  {'q_p99_ms':>10}  {'attach_ms':>10}")
    for s in summary:
        print(f"{s['M']:>2}  {s['ingest_wall_s']:>14.2f}  "
              f"{s['ingest_chunks_per_s']:>10.1f}  "
              f"{s['query_p50_ms']:>10.2f}  {s['query_p99_ms']:>10.2f}  "
              f"{s['query_attach_ms']:>10.2f}")
    print()
    print(f"results → {csv_path}")

    # Tidy up bench dir.
    if BENCH_DIR.exists():
        shutil.rmtree(BENCH_DIR)


if __name__ == "__main__":
    main()