LabDataStorageEvaluation/bench_runner.py

"""Task 09 - Benchmark execution and resource measurement.

Runs the 35-cell matrix (Q1-Q7 x 5 formats) in randomized order. Every
cell: 1 warm-up run (recorded as run 0, excluded from statistics) + 3
measured runs, each in its OWN subprocess sampled with psutil every 50 ms
for peak RSS / CPU time / read bytes. Result rows of every run are
validated against out/bench/expected/ (tolerance 1e-9) and recorded as
checksum_ok. Writes results_raw.csv, results_median.csv, PG EXPLAIN plans,
PG host/server metrics, the archival storage sizes, an environment record,
and the 09.ok marker. Spec: docs/specs/09_benchmark_execution.md; protocol:
docs/rules/bench-methodology.md.

Environment adaptations (bench-methodology section 6):
- the cold-cache matrix pass is SKIPPED (Windows host, no page-cache drop);
- PostgreSQL runs on a remote host: no backend-PID psutil sampling;
  instead best-effort Prometheus host CPU/RAM windows per PG cell and per
  session, plus best-effort pg_stat_statements per-query session deltas;
- on Windows the 50 ms sample reads the OS-maintained peak working set
  (peak_wset), so sub-interval children still report a true peak RSS.

After a run times out, the cell's remaining runs are SKIPPED (never retry
a timeout - bench-methodology section 3); only executed runs appear in the
raw results and the cell reports TIMEOUT.
"""

from __future__ import annotations

import csv
import gzip
import logging
import random
import re
import shutil
import statistics
import subprocess
import sys
import tarfile
import time
from dataclasses import dataclass, field
from pathlib import Path

import psutil
import psycopg

from common.monitoring import get_prom_url, host_window_metrics
from common.pg import get_dsn
from common.pipeline import (
	REPO_ROOT,
	check_dependencies,
	load_lab_config,
	parse_task_args,
	remove_stale_marker,
	write_marker,
)
from common.results import compare, parse_rows
from common.storage_sizes import update_storage_sizes

logger = logging.getLogger(__name__)

TASK_ID = "09"
DEPENDS_ON = ["08"]

CELL_TIMEOUT_S = 30 * 60   # hard per-run timeout, spec 09
SAMPLE_INTERVAL_S = 0.05   # psutil sampling cadence, spec 09
MEASURED_RUNS = 3          # 1 warm-up + 3 measured per cell, spec 09

FORMATS = ("csv", "json", "sqlite", "pg", "rdf")
QUERY_NUMBERS = tuple(range(1, 8))

RAW_HEADER = [
	"format", "query", "run", "wall_s", "peak_rss_mb", "cpu_s",
	"cpu_util_pct", "read_mb", "rows", "checksum_ok", "status", "exec_order",
]
MEDIAN_HEADER = [
	"format", "query", "wall_s_median", "wall_s_min", "wall_s_max",
	"peak_rss_mb_median", "cpu_s_median", "cpu_util_pct_median",
	"read_mb_median", "rows", "checksum_ok", "status",
]

# storage_sizes.csv completeness gate: every (format, variant) that must
# exist after consolidation (working + archival footprints, spec 09)
REQUIRED_SIZE_ROWS = {
	("csv", "tree"), ("csv", "targz"),
	("json", "full"), ("json", "hybrid"), ("json", "targz"),
	("sqlite", "db"), ("sqlite", "db_gz"),
	("pg", "live"), ("pg", "dump"),
	("rdf", "nt_gz"), ("rdf", "ttl"), ("rdf", "store"),
}


class CellTimeoutError(Exception):
	pass


@dataclass
class RunResult:
	status: str
	wall_s: float = 0.0
	peak_rss_mb: float = 0.0
	cpu_s: float = 0.0
	cpu_util_pct: float = 0.0
	read_mb: float = 0.0
	rows: int | None = None
	checksum_ok: bool | None = None

	def raw_row(self, fmt: str, qn: int, run: int, order: int) -> list:
		ok = "" if self.checksum_ok is None else str(self.checksum_ok).lower()
		return [
			fmt, f"q{qn}", run, round(self.wall_s, 4), round(self.peak_rss_mb, 1),
			round(self.cpu_s, 3), round(self.cpu_util_pct, 1), round(self.read_mb, 2),
			"" if self.rows is None else self.rows, ok, self.status, order,
		]


def _kill_tree(proc: subprocess.Popen) -> None:
	"""Kill the child and every descendant (the stub's real interpreter)."""
	try:
		for child in psutil.Process(proc.pid).children(recursive=True):
			try:
				child.kill()
			except psutil.NoSuchProcess:
				logger.debug("kill_tree: pid %d already exited", child.pid)
	except psutil.Error:
		logger.warning("kill_tree: descendants of pid %d not enumerable", proc.pid, exc_info=True)
	proc.kill()
	proc.wait()


def _run_child(cmd: list[str], stdout_path: Path, stderr_path: Path) -> tuple[int, RunResult]:
	"""Spawn one query subprocess and sample its process TREE until exit.

	The venv python.exe on Windows is a launcher stub: the real interpreter
	runs as its child, so sampling only the spawned PID measures the stub
	(~4 MB flat, ~0 CPU). Every sample therefore walks the descendant tree
	and keeps per-PID last-known CPU / read bytes and the OS peak working
	set; totals are sums over all PIDs ever seen. The flat stub adds ~4 MB
	to the peak-RSS sum - negligible against real query footprints.
	"""
	cores = psutil.cpu_count(logical=True)
	seen: dict[int, list] = {}  # pid -> [cpu_s, read_bytes, peak_rss_bytes]
	timed_out = False
	start = time.perf_counter()
	with open(stdout_path, "wb") as out_fh, open(stderr_path, "wb") as err_fh:
		proc = subprocess.Popen(cmd, stdout=out_fh, stderr=err_fh, cwd=REPO_ROOT)
		try:
			try:
				root = psutil.Process(proc.pid)
			except psutil.Error:
				root = None  # child already gone: exited faster than one sample
			while root is not None and proc.poll() is None:
				try:
					procs = [root, *root.children(recursive=True)]
				except psutil.Error:
					break
				for p in procs:
					try:
						mem = p.memory_info()
						cpu = p.cpu_times()
						io = p.io_counters()
					except psutil.Error:
						continue  # that process exited between listing and sampling
					entry = seen.setdefault(p.pid, [0.0, 0, 0])
					entry[0] = cpu.user + cpu.system
					entry[1] = io.read_bytes
					entry[2] = max(entry[2], getattr(mem, "peak_wset", 0) or mem.rss)
				if time.perf_counter() - start > CELL_TIMEOUT_S:
					timed_out = True
					break
				time.sleep(SAMPLE_INTERVAL_S)
			if timed_out:
				_kill_tree(proc)
				raise CellTimeoutError(f"exceeded {CELL_TIMEOUT_S} s")
			returncode = proc.wait()
		finally:
			if proc.poll() is None:
				_kill_tree(proc)
	wall = time.perf_counter() - start
	cpu_s = sum(entry[0] for entry in seen.values())
	result = RunResult(
		status="", wall_s=wall,
		peak_rss_mb=sum(entry[2] for entry in seen.values()) / 1048576,
		cpu_s=cpu_s,
		cpu_util_pct=100.0 * cpu_s / wall / cores if wall > 0 else 0.0,
		read_mb=sum(entry[1] for entry in seen.values()) / 1048576,
	)
	return returncode, result


def _run_cell_once(fmt: str, qn: int, cmd: list[str], tmp_dir: Path, expected: list[tuple]) -> RunResult:
	"""One benchmark run with pattern C isolation (code-error-handling)."""
	stdout_path = tmp_dir / "stdout.txt"
	stderr_path = tmp_dir / "stderr.txt"
	try:
		returncode, result = _run_child(cmd, stdout_path, stderr_path)
	except CellTimeoutError:
		logger.error("cell %s/q%d timed out after %d s", fmt, qn, CELL_TIMEOUT_S)
		return RunResult(status="TIMEOUT", wall_s=CELL_TIMEOUT_S)
	except Exception:
		logger.error("cell %s/q%d failed to execute", fmt, qn, exc_info=True)
		return RunResult(status="ERROR")
	if returncode != 0:
		tail = stderr_path.read_text(encoding="utf-8", errors="replace")[-2000:]
		logger.error("cell %s/q%d exited %d; stderr tail: %s", fmt, qn, returncode, tail)
		result.status = "ERROR"
		return result
	actual = parse_rows(stdout_path.read_text(encoding="ascii"))
	diff = compare(expected, actual)
	if diff:
		logger.error("cell %s/q%d checksum mismatch: %s", fmt, qn, diff)
	result.status = "OK"
	result.rows = len(actual)
	result.checksum_ok = diff is None
	return result


@dataclass
class CellOutcome:
	fmt: str
	qn: int
	measured: list[RunResult] = field(default_factory=list)
	timed_out: bool = False  # any run incl. warm-up: remaining runs skipped
	window: tuple[float, float] | None = None

	def status(self) -> str:
		if self.timed_out or any(r.status == "TIMEOUT" for r in self.measured):
			return "TIMEOUT"
		if len(self.measured) < MEASURED_RUNS or any(r.status == "ERROR" for r in self.measured):
			return "ERROR"
		if not all(r.checksum_ok for r in self.measured):
			return "INVALID"
		return "OK"


def _median_row(outcome: CellOutcome) -> list:
	status = outcome.status()
	fmt, qn, measured = outcome.fmt, outcome.qn, outcome.measured
	if status != "OK":
		ok = "false" if status == "INVALID" else ""
		return [fmt, f"q{qn}", "", "", "", "", "", "", "", "", ok, status]
	walls = [r.wall_s for r in measured]
	return [
		fmt, f"q{qn}",
		round(statistics.median(walls), 4), round(min(walls), 4), round(max(walls), 4),
		round(statistics.median(r.peak_rss_mb for r in measured), 1),
		round(statistics.median(r.cpu_s for r in measured), 3),
		round(statistics.median(r.cpu_util_pct for r in measured), 1),
		round(statistics.median(r.read_mb for r in measured), 2),
		measured[0].rows, "true", status,
	]


def _capture_explain(dsn: str, queries_dir: Path, explain_dir: Path) -> None:
	"""EXPLAIN (ANALYZE, BUFFERS) per PG query; verify Q5 avoids track_summary."""
	explain_dir.mkdir(parents=True, exist_ok=True)
	with psycopg.connect(dsn) as conn:
		for qn in QUERY_NUMBERS:
			sql = (queries_dir / "pg" / f"q{qn}.sql").read_text(encoding="ascii")
			plan = "\n".join(r[0] for r in conn.execute("EXPLAIN (ANALYZE, BUFFERS) " + sql).fetchall())
			(explain_dir / f"q{qn}.txt").write_text(plan + "\n", encoding="utf-8", newline="\n")
			if qn == 5 and "track_summary" in plan:
				raise RuntimeError("q5 plan touches track_summary - the forced full scan is violated (bench-methodology section 4)")
	logger.info("EXPLAIN plans captured to %s", explain_dir)


PG_STAT_QUERY = """
SELECT query, calls, total_exec_time, rows,
       shared_blks_hit, shared_blks_read, temp_blks_read, temp_blks_written
FROM pg_stat_statements
"""
PG_STAT_COLUMNS = [
	"calls", "total_exec_time_ms", "rows",
	"shared_blks_hit", "shared_blks_read", "temp_blks_read", "temp_blks_written",
]
_QTAG = re.compile(r"^-- Q(\d):")


def _pg_statements_snapshot(dsn: str) -> dict[int, list[float]] | None:
	"""Sum pg_stat_statements counters per benchmark query (matched by the
	leading '-- Q<N>:' comment). Best-effort: None when the extension is
	unavailable - server-side stats are auxiliary and never fail the task.
	"""
	try:
		sums: dict[int, list[float]] = {}
		with psycopg.connect(dsn) as conn:
			for query_text, *values in conn.execute(PG_STAT_QUERY):
				match = _QTAG.match(query_text)
				if not match:
					continue
				acc = sums.setdefault(int(match.group(1)), [0.0] * len(values))
				for i, v in enumerate(values):
					acc[i] += float(v)
		return sums
	except Exception:
		logger.warning("pg_stat_statements snapshot skipped (extension unavailable or not readable)", exc_info=True)
		return None


def _write_pg_server_stats(path: Path, start: dict[int, list[float]], end: dict[int, list[float]]) -> None:
	with open(path, "w", encoding="ascii", newline="") as fh:
		writer = csv.writer(fh)
		writer.writerow(["query"] + PG_STAT_COLUMNS)
		for qn in QUERY_NUMBERS:
			before = start.get(qn, [0.0] * len(PG_STAT_COLUMNS))
			after = end.get(qn)
			if after is None:
				logger.warning("pg_stat_statements: q%d missing in the end snapshot (entry evicted?)", qn)
				writer.writerow([f"q{qn}"] + [""] * len(PG_STAT_COLUMNS))
				continue
			delta = [after[i] - before[i] for i in range(len(PG_STAT_COLUMNS))]
			writer.writerow([f"q{qn}"] + [round(v, 3) for v in delta])
	logger.info("pg server-side statement deltas written to %s", path)


def _tree_bytes(root: Path) -> int:
	return sum(p.stat().st_size for p in root.rglob("*") if p.is_file())


def _consolidate_storage_sizes(out_root: Path, bench_dir: Path) -> None:
	"""Build archival variants, record all sizes, verify completeness."""
	archives = bench_dir / "archives"
	archives.mkdir(parents=True, exist_ok=True)

	csv_tgz = archives / "csv_tree.tar.gz"
	logger.info("archiving out/csv -> %s", csv_tgz)
	with tarfile.open(csv_tgz, "w:gz") as tar:
		tar.add(out_root / "csv", arcname="csv")
	update_storage_sizes(bench_dir, "csv", [
		("tree", _tree_bytes(out_root / "csv")), ("targz", csv_tgz.stat().st_size),
	])

	json_tgz = archives / "json.tar.gz"
	logger.info("archiving out/json -> %s", json_tgz)
	with tarfile.open(json_tgz, "w:gz") as tar:
		tar.add(out_root / "json", arcname="json")
	update_storage_sizes(bench_dir, "json", [("targz", json_tgz.stat().st_size)])

	db_gz = archives / "tribo.db.gz"
	logger.info("compressing out/sqlite/tribo.db -> %s", db_gz)
	with open(out_root / "sqlite" / "tribo.db", "rb") as src, gzip.open(db_gz, "wb") as dst:
		shutil.copyfileobj(src, dst, 1048576)
	update_storage_sizes(bench_dir, "sqlite", [("db_gz", db_gz.stat().st_size)])

	# pg_dump -Fc is internally gzip-compressed: it IS the pg archival variant
	update_storage_sizes(bench_dir, "pg", [("dump", (out_root / "pg" / "tribo.dump").stat().st_size)])

	present = set()
	with open(bench_dir / "storage_sizes.csv", encoding="ascii") as fh:
		for row in csv.reader(fh):
			present.add((row[0], row[1]))
	missing = REQUIRED_SIZE_ROWS - present
	if missing:
		raise ValueError(f"storage_sizes.csv incomplete after consolidation: missing {sorted(missing)}")
	logger.info("storage_sizes.csv consolidated: all %d required rows present", len(REQUIRED_SIZE_ROWS))


def _environment_record(path: Path, dsn: str) -> None:
	import platform

	lines = [
		f"benchmark_host_os: {platform.platform()}",
		f"benchmark_host_cpu: {platform.processor()}",
		f"benchmark_host_cores_logical: {psutil.cpu_count(logical=True)}",
		f"benchmark_host_cores_physical: {psutil.cpu_count(logical=False)}",
		f"benchmark_host_ram_gib: {round(psutil.virtual_memory().total / 1073741824, 1)}",
		f"python: {platform.python_version()}",
	]
	try:
		disks = subprocess.run(
			["powershell", "-NoProfile", "-Command",
				"Get-PhysicalDisk | ForEach-Object { $_.FriendlyName + ' / ' + $_.MediaType + ' / ' + [math]::Round($_.Size / 1GB) + ' GB' }"],
			capture_output=True, text=True, timeout=60, check=True)
		lines.append("benchmark_host_disks: " + "; ".join(ln.strip() for ln in disks.stdout.splitlines() if ln.strip()))
	except Exception:
		logger.warning("disk info skipped (Get-PhysicalDisk query failed - non-Windows host or no permission)", exc_info=True)
	import sqlite3

	import ijson
	import pyoxigraph

	lines += [
		f"sqlite: {sqlite3.sqlite_version}",
		f"psycopg: {psycopg.__version__}",
		f"pyoxigraph: {pyoxigraph.__version__}",
		f"ijson: {ijson.__version__}",
		f"psutil: {psutil.__version__}",
	]
	with psycopg.connect(dsn) as conn:
		pg_version = conn.execute("SELECT version()").fetchone()[0]
		lines.append(f"postgresql_server: {pg_version}")
		lines.append(f"postgresql_host: {conn.info.host} (remote; 4 cores / 3.8 GiB, see .done/06.ok)")
	lines += [
		"cold_cache_pass: skipped (Windows host, no page-cache drop - bench-methodology section 6)",
		"pg_backend_sampling: not possible for a remote server; Prometheus windows + pg_stat_statements deltas instead",
	]
	path.write_text("\n".join(lines) + "\n", encoding="utf-8", newline="\n")
	logger.info("environment record written to %s", path)


def _run_matrix(bench_dir: Path, expected: dict[int, list[tuple]], seed: int) -> list[CellOutcome]:
	cells = [(fmt, qn) for fmt in FORMATS for qn in QUERY_NUMBERS]
	random.Random(f"{seed}-bench-order").shuffle(cells)
	tmp_dir = bench_dir / "tmp"
	tmp_dir.mkdir(parents=True, exist_ok=True)
	outcomes: list[CellOutcome] = []
	with open(bench_dir / "results_raw.csv", "w", encoding="ascii", newline="") as raw_fh:
		writer = csv.writer(raw_fh)
		writer.writerow(RAW_HEADER)
		for order, (fmt, qn) in enumerate(cells, start=1):
			cmd = [sys.executable, str(bench_dir / "queries" / fmt / f"q{qn}.py")]
			outcome = CellOutcome(fmt, qn)
			window_start = time.time()
			for run in range(MEASURED_RUNS + 1):  # run 0 = warm-up
				result = _run_cell_once(fmt, qn, cmd, tmp_dir, expected[qn])
				writer.writerow(result.raw_row(fmt, qn, run, order))
				raw_fh.flush()
				if run > 0:
					outcome.measured.append(result)
				logger.info("cell %s/q%d run %d: %s wall=%.3fs rss=%.1fMB", fmt, qn, run, result.status, result.wall_s, result.peak_rss_mb)
				if result.status == "TIMEOUT":
					outcome.timed_out = True
					logger.error("cell %s/q%d: skipping remaining runs after timeout", fmt, qn)
					break
			outcome.window = (window_start, time.time())
			outcomes.append(outcome)
			status = outcome.status()
			walls = [r.wall_s for r in outcome.measured if r.status == "OK"]
			detail = f"median {statistics.median(walls):.3f} s" if status == "OK" else status
			print(f"[{order}/{len(cells)}] {fmt}/q{qn}: {detail}")
	shutil.rmtree(tmp_dir)
	return outcomes


def _write_pg_host_metrics(path: Path, prom_url: str, outcomes: list[CellOutcome]) -> None:
	with open(path, "w", encoding="ascii", newline="") as fh:
		writer = csv.writer(fh)
		writer.writerow([
			"query", "window_s", "host_cpu_avg_pct", "host_cpu_max_pct",
			"host_ram_used_baseline_mb", "host_ram_used_peak_mb", "host_ram_used_delta_mb",
		])
		for outcome in sorted((o for o in outcomes if o.fmt == "pg"), key=lambda o: o.qn):
			start, end = outcome.window
			metrics = host_window_metrics(prom_url, start, end)
			row = [f"q{outcome.qn}", round(end - start, 1)]
			if metrics is None:
				row += [""] * 5
			else:
				row += [
					metrics["host_cpu_avg_pct"], metrics["host_cpu_max_pct"],
					metrics["host_ram_used_baseline_mb"], metrics["host_ram_used_peak_mb"],
					metrics["host_ram_used_delta_mb"],
				]
			writer.writerow(row)
	logger.info("per-cell PG host metrics written to %s", path)


def main() -> int:
	args = parse_task_args("Task 09: execute the Q1-Q7 x 5 formats benchmark matrix")
	out_root: Path = args.out_root
	bench_dir = out_root / "bench"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		dump_path = out_root / "pg" / "tribo.dump"
		if not dump_path.exists():
			raise FileNotFoundError(f"{dump_path} missing - produce it with pg_dump -Fc before task 09 (storage consolidation needs it)")
		dsn = get_dsn()
		expected = {qn: parse_rows((bench_dir / "expected" / f"q{qn}.rows.csv").read_text(encoding="ascii")) for qn in QUERY_NUMBERS}
		seed = load_lab_config(out_root)["seed"]
		logger.warning("cold-cache pass skipped (Windows host - bench-methodology section 6)")
		_environment_record(bench_dir / "environment.txt", dsn)

		prom_url = get_prom_url()
		stats_start = _pg_statements_snapshot(dsn)
		session_start = time.time()
		outcomes = _run_matrix(bench_dir, expected, seed)
		session_end = time.time()

		with open(bench_dir / "results_median.csv", "w", encoding="ascii", newline="") as fh:
			writer = csv.writer(fh)
			writer.writerow(MEDIAN_HEADER)
			for outcome in sorted(outcomes, key=lambda o: (o.fmt, o.qn)):
				writer.writerow(_median_row(outcome))

		# snapshot BEFORE the EXPLAIN pass so its executions cannot pollute the deltas
		if stats_start is not None:
			stats_end = _pg_statements_snapshot(dsn)
			if stats_end is not None:
				_write_pg_server_stats(bench_dir / "pg_server_stats.csv", stats_start, stats_end)
		_capture_explain(dsn, bench_dir / "queries", bench_dir / "explain")
		session_host = None
		if prom_url:
			_write_pg_host_metrics(bench_dir / "pg_host_metrics.csv", prom_url, outcomes)
			session_host = host_window_metrics(prom_url, session_start, session_end)
		else:
			logger.warning("TRIBO_PROM_URL unset - PG host metrics skipped")

		_consolidate_storage_sizes(out_root, bench_dir)

		by_status: dict[str, int] = {}
		for outcome in outcomes:
			by_status[outcome.status()] = by_status.get(outcome.status(), 0) + 1
		entries = {
			"cells_total": len(outcomes),
			"cells_ok": by_status.get("OK", 0),
			"cells_invalid": by_status.get("INVALID", 0),
			"cells_timeout": by_status.get("TIMEOUT", 0),
			"cells_error": by_status.get("ERROR", 0),
			"session_wall_s": round(session_end - session_start, 1),
			"cold_cache_pass": "skipped (Windows host)",
		}
		if session_host:
			entries.update({f"session_{k}": v for k, v in session_host.items()})
		marker_path = write_marker(out_root, TASK_ID, entries)
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	print(f"task 09 ok: {entries['cells_ok']}/{entries['cells_total']} cells OK "
		f"(invalid={entries['cells_invalid']}, timeout={entries['cells_timeout']}, error={entries['cells_error']}), "
		f"session {entries['session_wall_s']} s")
	print(f"results: {bench_dir / 'results_raw.csv'}, {bench_dir / 'results_median.csv'}")
	print(f"marker: {marker_path}")
	return 0


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