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 shutil
import statistics
import tempfile
import time
from pathlib import Path

from codexa.store.generations import (
    GenerationCloneStats,
    begin_generation_build,
    publish_generation_if_staged,
)
from codexa.indexing.progressive_ready_proof import ReadyProof
from codexa.store.progressive_generation_publish import (
    commit_ready_generation,
    mark_generation_ready,
)
from codexa.store.progressive_generations import _progressive_clone_base

_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 _progressive_proof(generation: str) -> ReadyProof:
    return ReadyProof(
        build_id="benchmark", generation=generation, checkpoint_seq=1,
        base_generation=None, checkpoint_mode="delta", contract="benchmark",
        revision="benchmark", target_digest="a" * 64,
        manifest_digest="b" * 64, manifest_records_digest="c" * 64,
        dense_digest="d" * 64, manifest_sources=1, manifest_chunks=1,
        dense_sources=1, dense_chunks=1, bm25_enabled=False,
        bm25_digest=None, bm25_sources=None, bm25_rows=None,
        colbert_enabled=False, colbert_digest=None, colbert_sources=None,
        colbert_rows=None, bm25_inventory_digest=None,
        colbert_inventory_digest=None, index_info_digest="e" * 64,
        indexed=1, total=2, partial_corpus=True, processed_ordinals=(0,),
        indexed_ordinals=(0,), dense_inventory_digest="f" * 64,
        batch_mapping_digest="1" * 64, cumulative_source_digest="2" * 64,
        committed_phase_digest="3" * 64, phase_state_digest="4" * 64,
        phase_transition_digest="5" * 64,
    )


def _one_progressive_run(
    buckets: int,
) -> tuple[GenerationCloneStats, float]:
    with tempfile.TemporaryDirectory(prefix="codexa-rob116-") as temp:
        root = Path(temp) / "clone"
        _seed(root, buckets)
        staging = root / "gen_000002"
        stats = _progressive_clone_base(
            root,
            staging,
            build_id="benchmark",
            checkpoint_seq=2,
            base_generation="gen_000001",
            lineage_fields={
                "lineage_digest": "0" * 64,
                "base_checkpoint_seq": 1,
            },
            manifest_name="manifest.json",
            index_info_name="index_info.json",
        )
        if stats is None:
            raise RuntimeError("progressive clone did not run")
        shutil.rmtree(staging)

        publish_root = Path(temp) / "publish"
        cfg = _cfg(publish_root)
        build = begin_generation_build(
            cfg,
            changed_sources=[],
            force_reindex=True,
            manifest_path=str(publish_root / "manifest.json"),
            index_info_path=str(publish_root / "index_info.json"),
            log=logging.getLogger("codexa.rob116.bench"),
        )
        if build is None:
            raise RuntimeError("progressive publication stage did not open")
        proof = _progressive_proof(build.generation_id)
        mark_generation_ready(
            build.cfg,
            proof_json=proof.as_json(),
            base_generation=None,
            log=logging.getLogger("codexa.rob116.bench"),
        )
        started = time.perf_counter()
        published = commit_ready_generation(
            build.cfg,
            proof_json=proof.as_json(),
            base_generation=None,
            validate=lambda: proof,
            log=logging.getLogger("codexa.rob116.bench"),
        )
        pointer_ms = (time.perf_counter() - started) * 1000.0
        if not published:
            raise RuntimeError("progressive pointer publication did not complete")
        return stats, pointer_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)]
    progressive = [
        _one_progressive_run(args.buckets) for _ in range(args.runs)
    ]
    stats = samples[-1][0]
    progressive_stats = progressive[-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
        ),
        "progressive_hardlinked_files": progressive_stats.hardlinked_files,
        "progressive_copied_files": progressive_stats.copied_files,
        "progressive_hardlink_fallback_files": (
            progressive_stats.hardlink_fallback_files
        ),
        "median_progressive_clone_ms": statistics.median(
            sample.clone_duration_s * 1000.0 for sample, _ in progressive
        ),
        "median_progressive_pointer_ms": statistics.median(
            pointer_ms for _, pointer_ms in progressive
        ),
    }
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()