arborist/bench/cold_pack_roundtrip.py

"""Producer/consumer benchmark for the cold-pack distribution tier
(#000061) running against a real DO Spaces bucket and a real fresh
peer on `3090-ai.foxhop.net`.

Records timings + byte counts in both directions:

  PRODUCER  (this host)
    pack wall time, compressed MB/s, packs created, bucket bytes
  CONSUMER  (3090)
    hydrate wall time, downloaded MB/s, rows restored per table,
    shard sizes on disk, post-hydrate row count match

Result lands at ``bench/results/cold-pack-roundtrip-<ISO timestamp>.json``
so a future operator (or follow-on ticket re-running this benchmark
after pack-format changes) has a single artifact to diff against.

Run after pack has been produced + bucket has packs to pull. See the
`if __name__ == "__main__"` driver for the exact sequence.

Read-only on the producer host (just queries process state + bucket
stats). On the consumer host: SSH'd command runs `arborist cold
unpack --hydrate-shards-dir ~/.arborist/shards --hydrate-M 4 --full`.
"""
from __future__ import annotations

import json
import os
import shlex
import subprocess
import time
from pathlib import Path

RESULTS_DIR = Path(__file__).resolve().parent / "results"
SHARDS_DIR = Path("/home/fox/.arborist/shards")
COLD_ENV = {
    "ARBORIST_COLD_ENDPOINT_URL": "https://nyc3.digitaloceanspaces.com",
    "ARBORIST_COLD_BUCKET": "arborist",
}
ARBORIST = "/home/fox/git/arborist/.venv/bin/arborist"
CONSUMER_HOST = "3090-ai.foxhop.net"
CONSUMER_ARBORIST = "/home/fox/git/arborist/.venv/bin/arborist"


def _cold_stats_local() -> dict:
    """Call `arborist cold stats` on this host; return parsed JSON."""
    env = {**os.environ, **COLD_ENV}
    out = subprocess.check_output(
        [ARBORIST, "cold", "stats"], env=env, text=True
    )
    return json.loads(out)


def _cold_list_local() -> dict:
    env = {**os.environ, **COLD_ENV}
    out = subprocess.check_output(
        [ARBORIST, "cold", "list", "--no-manifest"], env=env, text=True
    )
    return json.loads(out)


def _identify_metadata_pack_hash() -> str:
    """Read manifest/latest.json from the bucket to find the active
    metadata pack hash. Fresh peer would do the same on genesis."""
    env = {**os.environ, **COLD_ENV}
    # Use `arborist cold list` (it sorts manifest packs first).
    listing = _cold_list_local()
    for entry in listing["packs"]:
        if entry.get("kind") == "metadata":
            return entry["pack_hash"]
    raise RuntimeError("no metadata pack found in bucket")


def _ssh(cmd: str) -> subprocess.CompletedProcess:
    """Run a shell command on the consumer host. Returns CompletedProcess."""
    return subprocess.run(
        ["ssh", "-o", "BatchMode=yes", CONSUMER_HOST, cmd],
        capture_output=True, text=True,
    )


def measure_consumer(metadata_pack_hash: str, *, M: int = 4) -> dict:
    """Run the hydrate-into-empty path on 3090, time it, return metrics."""
    started_at = time.time()
    cmd = " && ".join([
        # Belt-and-suspenders: ensure no stale shards exist.
        "rm -rf ~/.arborist/shards-genesis-test",
        "mkdir -p ~/.arborist/shards-genesis-test",
        # Run hydrate; pipe output to a side log so we can diff later.
        "/usr/bin/time -v "
        f"{CONSUMER_ARBORIST} cold unpack {shlex.quote(metadata_pack_hash)} "
        f"--hydrate-shards-dir ~/.arborist/shards-genesis-test "
        f"--hydrate-M {M} --full "
        "2>&1 | tee /tmp/cold-genesis.log",
    ])
    result = _ssh(cmd)
    elapsed = time.time() - started_at
    # Capture final shard sizes from the consumer.
    sizes_raw = _ssh(
        "ls -la ~/.arborist/shards-genesis-test/*.db 2>/dev/null | "
        "awk '{print $NF, $5}'"
    ).stdout
    shard_sizes: dict[str, int] = {}
    for line in sizes_raw.strip().splitlines():
        if not line.strip():
            continue
        parts = line.split()
        shard_sizes[Path(parts[0]).name] = int(parts[1])
    # And row counts per shard.
    row_counts: dict[str, dict[str, int]] = {}
    for i in range(M):
        cmd = (
            "python3 -c \"import sqlite3; "
            "c=sqlite3.connect('~/.arborist/shards-genesis-test/{:03d}.db'.expanduser()) "
            "if hasattr(str,'expanduser') else "
            "sqlite3.connect('/home/fox/.arborist/shards-genesis-test/{:03d}.db'); "
            "print(c.execute('SELECT COUNT(*) FROM documents').fetchone()[0], "
            "c.execute('SELECT COUNT(*) FROM chunks').fetchone()[0], "
            "c.execute('SELECT COUNT(*) FROM edges').fetchone()[0])\""
        ).format(i, i)
        r = _ssh(cmd)
        parts = r.stdout.strip().split()
        if len(parts) >= 3:
            row_counts[f"{i:03d}.db"] = {
                "documents": int(parts[0]),
                "chunks": int(parts[1]),
                "edges": int(parts[2]),
            }
    return {
        "elapsed_seconds": elapsed,
        "ssh_exit": result.returncode,
        "stdout_tail": result.stdout[-2000:] if result.stdout else "",
        "stderr_tail": result.stderr[-2000:] if result.stderr else "",
        "shard_sizes_bytes": shard_sizes,
        "row_counts_by_shard": row_counts,
    }


def main(producer_started_at: float | None = None) -> int:
    """Driver: read producer-side stats, run consumer benchmark, save JSON."""
    finished_at = time.time()
    bucket = _cold_stats_local()
    pack_hash = _identify_metadata_pack_hash()
    print(f"metadata pack: {pack_hash[:16]}…")
    print(f"bucket: {bucket['packs']} packs / "
          f"{bucket['packs_compressed_bytes'] / 1e9:.2f} GB")
    print(f"running consumer on {CONSUMER_HOST} …")
    consumer = measure_consumer(pack_hash, M=4)

    out = {
        "captured_at_utc": time.strftime(
            "%Y-%m-%dT%H:%M:%SZ", time.gmtime()
        ),
        "producer_host": os.uname().nodename,
        "consumer_host": CONSUMER_HOST,
        "metadata_pack_hash": pack_hash,
        "bucket": bucket,
        "consumer": consumer,
    }
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    out_path = (
        RESULTS_DIR
        / f"cold-pack-roundtrip-{time.strftime('%Y-%m-%dT%H-%M-%SZ', time.gmtime())}.json"
    )
    out_path.write_text(json.dumps(out, indent=2))
    print(f"wrote {out_path}")
    print(f"consumer elapsed: {consumer['elapsed_seconds']:.1f}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())