bench_generation_clone.py

#!/usr/bin/env python3
"""Reproduce SD-12 generation clone/publish bookkeeping evidence."""

from __future__ import annotations

import argparse
import json
import logging
import statistics
import tempfile
import time
from pathlib import Path

from codexa.store.generations import (
    GenerationCloneStats,
    begin_generation_build,
    publish_generation_if_staged,
)

_MIB = 1024 * 1024


def _cfg(root: Path) -> dict:
    return {
        "store": {
            "backend": "sharded_chroma",
            "persist_dir": str(root),
            "publish": "generations",
            "retain_generations": 2,
        },
        "metadata": {
            "manifest_file": str(root / "manifest.json"),
            "index_info_file": str(root / "index_info.json"),
        },
    }


def _sized_file(path: Path, size: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("wb") as stream:
        stream.truncate(size)


def _seed(root: Path, buckets: int) -> None:
    current = root / "gen_000001"
    current.mkdir(parents=True)
    for index in range(buckets):
        stem = "txt_a_0000000000" if index == 0 else f"pdf_b_{index:010d}"
        _sized_file(current / f"{stem}.sqlite3", _MIB)
        shard = current / stem
        _sized_file(shard / "segment.bin", _MIB // 2)
        (shard / "chroma.sqlite3").symlink_to(Path("..") / f"{stem}.sqlite3")
    _sized_file(current / "manifest.json", 64 * 1024)
    (current / "index_info.json").write_text("{}", encoding="utf-8")
    _sized_file(current / "bm25.db", _MIB // 2)
    _sized_file(current / "colbert_doc_embeddings.sqlite3", _MIB // 2)
    (root / "CURRENT").write_text("gen_000001\n", encoding="utf-8")


def _one_run(buckets: int) -> tuple[GenerationCloneStats, float]:
    with tempfile.TemporaryDirectory(prefix="codexa-sd12-") as temp:
        root = Path(temp) / "store"
        _seed(root, buckets)
        cfg = _cfg(root)
        log = logging.getLogger("codexa.sd12.bench")
        build = begin_generation_build(
            cfg,
            changed_sources=["/data/apple.txt"],
            force_reindex=False,
            manifest_path=str(root / "manifest.json"),
            index_info_path=str(root / "index_info.json"),
            log=log,
        )
        if build is None or build.clone_stats is None:
            raise RuntimeError("generation clone did not run")
        started = time.perf_counter()
        publish_generation_if_staged(build.cfg, log)
        publish_ms = (time.perf_counter() - started) * 1000.0
        return build.clone_stats, publish_ms


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=5)
    parser.add_argument("--buckets", type=int, default=32)
    args = parser.parse_args()
    if args.runs < 1 or args.buckets < 2:
        parser.error("--runs must be >= 1 and --buckets must be >= 2")
    samples = [_one_run(args.buckets) for _ in range(args.runs)]
    stats = samples[-1][0]
    payload = {
        "runs": args.runs,
        "buckets": args.buckets,
        "logical_bytes": stats.logical_bytes,
        "copied_files": stats.copied_files,
        "copied_bytes": stats.copied_bytes,
        "copied_ratio": stats.copied_bytes / stats.logical_bytes,
        "hardlinked_files": stats.hardlinked_files,
        "hardlinked_bytes": stats.hardlinked_bytes,
        "hardlink_fallback_files": stats.hardlink_fallback_files,
        "median_clone_ms": statistics.median(
            sample.clone_duration_s * 1000.0 for sample, _ in samples
        ),
        "median_publish_gc_ms": statistics.median(
            publish_ms for _, publish_ms in samples
        ),
    }
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()