arborist/scripts/backfill_concepts.py

"""Run every concept-layer extractor across every shard, in parallel.

Replaces the ad-hoc `python -c "..."` invocations the README + design
doc currently document. Wall-clock parallelism comes from a process
pool over the shard set — each shard's work is self-contained (separate
SQLite file, separate FTS5 index), so per-shard processes don't
contend.

Wired into `make backfill-concepts`. Usage:

    python scripts/backfill_concepts.py --shards-dir ~/.aborist/shards
    python scripts/backfill_concepts.py --shards-dir ~/.aborist/shards \
        --extractors link_reciprocity,token_idf,documents_fts
    python scripts/backfill_concepts.py --shards-dir ~/.aborist/shards \
        --workers 8

Skips shards that aren't numeric-stem (`000.db`, `001.db`, …) by
default — the qa.db / snapshots.db / crawl_*.db shards rarely need
concept-layer backfill, and including them confuses progress reporting.
Pass --include-non-numeric to opt in.
"""

from __future__ import annotations

import argparse
import multiprocessing as mp
import sys
import time
from pathlib import Path

from aborist.concepts.extract import EXTRACTORS
from aborist.store import connect, discover_shards


# Default extractor order. link_reciprocity must run BEFORE token_idf
# (idf indexes the token set that link_reciprocity populated). Order
# matters when a shard has no concept_relations rows yet.
DEFAULT_EXTRACTORS = ("link_reciprocity", "token_idf", "documents_fts")


def _run_one_shard(args: tuple[Path, tuple[str, ...]]) -> dict:
    shard, extractors = args
    name = shard.stem
    t0 = time.time()
    results: dict[str, dict] = {}
    conn = connect(shard)
    try:
        for kind in extractors:
            fn = EXTRACTORS.get(kind)
            if fn is None:
                results[kind] = {"error": f"unknown extractor {kind!r}"}
                continue
            results[kind] = fn(conn, derived_from=f"backfill@{name}")
    finally:
        conn.close()
    return {
        "shard": str(shard),
        "name": name,
        "elapsed_s": round(time.time() - t0, 2),
        "results": results,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.strip().splitlines()[0])
    parser.add_argument(
        "--shards-dir",
        type=Path,
        default=Path.home() / ".aborist" / "shards",
        help="Directory containing *.db shards.",
    )
    parser.add_argument(
        "--extractors",
        default=",".join(DEFAULT_EXTRACTORS),
        help=(
            "Comma-separated extractor names (registered in "
            "aborist.concepts.extract.EXTRACTORS). Default: "
            f"{','.join(DEFAULT_EXTRACTORS)}"
        ),
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Parallel shard workers. Default 4.",
    )
    parser.add_argument(
        "--include-non-numeric",
        action="store_true",
        help=(
            "Include shards whose stem isn't a digit run (qa.db, "
            "snapshots.db, crawl_*.db). Off by default."
        ),
    )
    ns = parser.parse_args(argv)

    extractors = tuple(e.strip() for e in ns.extractors.split(",") if e.strip())
    for e in extractors:
        if e not in EXTRACTORS:
            parser.error(
                f"unknown extractor {e!r}; available: "
                f"{sorted(EXTRACTORS)}"
            )

    shards = [
        s for s in discover_shards(ns.shards_dir)
        if ns.include_non_numeric or s.stem.isdigit()
    ]
    if not shards:
        parser.error(f"no shards found under {ns.shards_dir}")

    print(
        f">> {len(shards)} shard(s) × {len(extractors)} extractor(s) "
        f"× {ns.workers} workers",
        flush=True,
    )

    grand_t0 = time.time()
    work = [(s, extractors) for s in shards]
    with mp.Pool(processes=ns.workers) as pool:
        for outcome in pool.imap_unordered(_run_one_shard, work):
            print(f"[{outcome['name']}] {outcome['elapsed_s']}s", flush=True)
            for kind, result in outcome["results"].items():
                print(f"   {kind}: {result}", flush=True)

    print(f"\n>> total wall-clock: {round(time.time() - grand_t0, 1)}s")
    return 0


if __name__ == "__main__":
    sys.exit(main())