From d942685525e36d8ece287e4061dfba012198dc47 Mon Sep 17 00:00:00 2001 From: Michael C Li Date: Mon, 6 Jul 2026 18:37:21 +0100 Subject: [PATCH] Add window1 P3 representative run orchestration --- .../testbench/window1_p3/__init__.py | 1 + .../testbench/window1_p3/run_p3_rev_subset.py | 1551 +++++++++++++++++ .../window1_p3/run_p3_scale_batch.py | 722 ++++++++ .../window1_p3/wait_p3_rev_stage_b.py | 288 +++ .../window1_p3/watch_p3_rev_raw_build.py | 136 ++ 5 files changed, 2698 insertions(+) create mode 100644 transactional_emulator/testbench/window1_p3/__init__.py create mode 100644 transactional_emulator/testbench/window1_p3/run_p3_rev_subset.py create mode 100644 transactional_emulator/testbench/window1_p3/run_p3_scale_batch.py create mode 100644 transactional_emulator/testbench/window1_p3/wait_p3_rev_stage_b.py create mode 100644 transactional_emulator/testbench/window1_p3/watch_p3_rev_raw_build.py diff --git a/transactional_emulator/testbench/window1_p3/__init__.py b/transactional_emulator/testbench/window1_p3/__init__.py new file mode 100644 index 00000000..f8eef558 --- /dev/null +++ b/transactional_emulator/testbench/window1_p3/__init__.py @@ -0,0 +1 @@ +"""Window 1 P3 representative and scale experiment orchestration.""" diff --git a/transactional_emulator/testbench/window1_p3/run_p3_rev_subset.py b/transactional_emulator/testbench/window1_p3/run_p3_rev_subset.py new file mode 100644 index 00000000..952d11d3 --- /dev/null +++ b/transactional_emulator/testbench/window1_p3/run_p3_rev_subset.py @@ -0,0 +1,1551 @@ +#!/usr/bin/env python3 +"""Window 1 P3-rev representative subset and parallel replay runner. + +This sidecar keeps the P2 replay semantics intact. It only changes scale +strategy: select a reproducible representative subset, generate true routes for +missing samples, then replay selected route traces in parallel with checkpointed +progress. +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import csv +import hashlib +import json +import os +import random +import re +import shutil +import statistics +import subprocess +import sys +import time +import traceback +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import PLENA_ROOT +from transactional_emulator.testbench.window1_p2 import ( + build_route_traces, + export_pilot_results, + generate_true_routing_with_weights, +) +from transactional_emulator.testbench.window1_p2.p2_utils import ( + INPUT_FILES, + MODEL_CONFIGS, + ensure_paths, + iter_jsonl, + load_json, + write_csv, + write_json, +) +from transactional_emulator.testbench.window1_p2.run_trace_batch import _prune_success_artifacts + + +DEFAULT_OUT_ROOT = PLENA_ROOT / "outputs" / "window1_p3" +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_LAYERS = [0, 12, 23] +DEFAULT_SEED = 20260704 +MEASUREMENT_NOTE = "self-consistent upper bound, absolute accuracy pending RTL (Window 2)" +ROUTING_STAGE_NAMES = { + "accumulator_init", + "gather", + "expert_weight_address", + "expert_route_weight", + "scatter_combine", +} + + +def _parse_csv_strings(value: str | None) -> list[str]: + if not value: + return [] + return [part.strip() for part in value.split(",") if part.strip()] + + +def _parse_csv_ints(value: str | None) -> list[int]: + if not value: + return [] + return [int(part) for part in value.split(",") if part.strip()] + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f"{path.name}.tmp.{os.getpid()}") + tmp.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + tmp.replace(path) + + +def _append_jsonl(path: Path, row: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _stable_seed(*parts: object) -> int: + text = ":".join(str(part) for part in parts) + digest = hashlib.sha256(text.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") + + +def _stats(values: list[int | float | None]) -> dict[str, Any]: + clean = [float(value) for value in values if value is not None] + if not clean: + return {"count": 0, "min": None, "median": None, "mean": None, "p95": None, "max": None} + ordered = sorted(clean) + if len(ordered) == 1: + p95 = ordered[0] + else: + rank = (len(ordered) - 1) * 0.95 + lo = int(rank) + hi = min(lo + 1, len(ordered) - 1) + frac = rank - lo + p95 = ordered[lo] * (1.0 - frac) + ordered[hi] * frac + return { + "count": len(clean), + "min": min(clean), + "median": float(statistics.median(clean)), + "mean": float(statistics.mean(clean)), + "p95": float(p95), + "max": max(clean), + } + + +def _fmt_seconds(value: float | int | None) -> str: + if value is None: + return "n/a" + value_f = float(value) + return f"{value_f:.1f}s ({value_f / 3600.0:.2f}h)" + + +def _fmt_bytes(value: float | int | None) -> str: + if value is None: + return "n/a" + value_f = float(value) + units = ["B", "KiB", "MiB", "GiB", "TiB"] + idx = 0 + while value_f >= 1024.0 and idx < len(units) - 1: + value_f /= 1024.0 + idx += 1 + return f"{value_f:.2f} {units[idx]}" + + +def _input_records(benchmark: str, model_key: str) -> list[dict[str, Any]]: + rows = [] + for idx, row in enumerate(iter_jsonl(INPUT_FILES[benchmark])): + if row.get("model_key") != model_key: + continue + rows.append( + { + "benchmark": benchmark, + "model_key": model_key, + "sample_id": str(row["sample_id"]), + "sample_index": idx, + "category": str(row.get("category") or "uncategorized"), + "input_tokens": int(row.get("input_tokens") or len(row.get("input_ids") or [])), + "input_chars": row.get("input_chars"), + "source_file": row.get("source_file"), + } + ) + return rows + + +def _largest_remainder_quotas(counts: Counter[str], target: int) -> dict[str, int]: + if target <= 0: + return {} + total = sum(counts.values()) + if total <= 0: + return {} + categories = sorted(counts) + raw = {category: (counts[category] * float(target)) / float(total) for category in categories} + quotas = {category: min(counts[category], int(raw[category])) for category in categories} + for category in categories: + if counts[category] > 0 and quotas[category] == 0 and target >= len(categories): + quotas[category] = 1 + while sum(quotas.values()) > target: + candidates = [category for category in categories if quotas[category] > 1] + category = min(candidates, key=lambda item: (raw[item] - int(raw[item]), quotas[item], item)) + quotas[category] -= 1 + while sum(quotas.values()) < target: + candidates = [category for category in categories if quotas[category] < counts[category]] + if not candidates: + break + category = max(candidates, key=lambda item: (raw[item] - int(raw[item]), counts[item], item)) + quotas[category] += 1 + return quotas + + +def _select_bfcl(rows: list[dict[str, Any]], target: int, seed: int) -> dict[str, Any]: + by_category: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + by_category[row["category"]].append(row) + counts = Counter({category: len(items) for category, items in by_category.items()}) + quotas = _largest_remainder_quotas(counts, min(target, len(rows))) + selected: list[dict[str, Any]] = [] + categories = [] + for category in sorted(by_category): + items = sorted(by_category[category], key=lambda row: (row["sample_index"], row["sample_id"])) + quota = int(quotas.get(category, 0)) + rng = random.Random(_stable_seed(seed, "bfcl_v3", category)) + shuffled = items[:] + rng.shuffle(shuffled) + chosen = sorted(shuffled[:quota], key=lambda row: (row["sample_index"], row["sample_id"])) + selected.extend(chosen) + categories.append( + { + "category": category, + "available": len(items), + "quota": quota, + "selected_sample_ids": [row["sample_id"] for row in chosen], + } + ) + selected.sort(key=lambda row: (row["sample_index"], row["sample_id"])) + return { + "available": len(rows), + "target": target, + "selected_count": len(selected), + "selection_method": "category_stratified_largest_remainder_fixed_seed", + "categories": categories, + "selected_samples": selected, + } + + +def _select_gpqa(rows: list[dict[str, Any]]) -> dict[str, Any]: + selected = sorted(rows, key=lambda row: (row["sample_index"], row["sample_id"])) + categories = [] + by_category: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in selected: + by_category[row["category"]].append(row) + for category in sorted(by_category): + categories.append( + { + "category": category, + "available": len(by_category[category]), + "quota": len(by_category[category]), + "selected_sample_ids": [row["sample_id"] for row in by_category[category]], + } + ) + return { + "available": len(rows), + "target": "all", + "selected_count": len(selected), + "selection_method": "all_samples_no_sampling", + "categories": categories, + "selected_samples": selected, + } + + +def select_subset(args: argparse.Namespace) -> dict[str, Any]: + payload = { + "schema_version": 1, + "created_at_unix": time.time(), + "seed": args.seed, + "model_key": args.model_key, + "model_name": MODEL_CONFIGS[args.model_key]["name"], + "layers": args.layers, + "phase": args.write_phases, + "decode_steps": args.decode_steps, + "measurement_note": MEASUREMENT_NOTE, + "bfcl_target": args.bfcl_target, + "benchmarks": {}, + } + if "bfcl_v3" in args.benchmarks: + payload["benchmarks"]["bfcl_v3"] = _select_bfcl( + _input_records("bfcl_v3", args.model_key), + args.bfcl_target, + args.seed, + ) + if "gpqa_diamond" in args.benchmarks: + payload["benchmarks"]["gpqa_diamond"] = _select_gpqa(_input_records("gpqa_diamond", args.model_key)) + write_json(args.selected_out, payload) + return payload + + +def _load_selection(args: argparse.Namespace) -> dict[str, Any]: + if args.selected_out.exists() and not args.reselect: + return load_json(args.selected_out) + return select_subset(args) + + +def _sample_ids(selection: dict[str, Any], benchmark: str) -> list[str]: + bench = selection["benchmarks"].get(benchmark, {}) + return [str(row["sample_id"]) for row in bench.get("selected_samples", [])] + + +def _sample_ids_for_args(selection: dict[str, Any], benchmark: str, args: argparse.Namespace) -> list[str]: + bench = selection["benchmarks"].get(benchmark, {}) + rows = list(bench.get("selected_samples", [])) + shard_count = getattr(args, "sample_shard_count", None) + shard_index = getattr(args, "sample_shard_index", None) + if shard_count is not None: + if shard_count <= 0: + raise ValueError("--sample-shard-count must be positive") + if shard_index is None or not (0 <= shard_index < shard_count): + raise ValueError("--sample-shard-index must be in [0, sample_shard_count)") + balanced = sorted( + rows, + key=lambda row: (-int(row.get("input_tokens") or 0), str(row.get("sample_id"))), + ) + rows = [row for idx, row in enumerate(balanced) if idx % shard_count == shard_index] + rows.sort(key=lambda row: (int(row.get("sample_index") or 0), str(row.get("sample_id")))) + return [str(row["sample_id"]) for row in rows] + + +def _routing_output(out_root: Path, benchmark: str, model_key: str, layers: list[int], args: argparse.Namespace) -> Path: + layer_text = "_".join(str(layer) for layer in layers) + suffix = str(getattr(args, "routing_output_suffix", "") or "") + return ( + out_root + / "true_routing" + / f"{benchmark}_{model_key}_layers_{layer_text}_{args.write_phases}_d{args.decode_steps}{suffix}.jsonl" + ) + + +def _routing_resume_sources( + out_root: Path, + benchmark: str, + model_key: str, + layers: list[int], + args: argparse.Namespace, +) -> list[Path]: + layer_text = "_".join(str(layer) for layer in layers) + return sorted( + (out_root / "true_routing").glob( + f"{benchmark}_{model_key}_layers_{layer_text}_{args.write_phases}_d{args.decode_steps}*.jsonl" + ) + ) + + +def run_route_generation(args: argparse.Namespace, selection: dict[str, Any]) -> list[dict[str, Any]]: + summaries = [] + log_suffix = str(getattr(args, "routing_log_suffix", "") or "") + for benchmark in args.benchmarks: + ids = _sample_ids_for_args(selection, benchmark, args) + if not ids: + continue + routing_args = argparse.Namespace( + benchmark=benchmark, + model_key=args.model_key, + limit=None, + sample_ids=",".join(ids), + max_input_tokens=args.max_input_tokens, + layers=args.layers, + decode_steps=args.decode_steps, + threads=args.threads, + batch_size=args.batch_size, + output=_routing_output(args.out_root, benchmark, args.model_key, args.layers, args), + resume=True, + sort_by_length=args.sort_by_length, + write_phases=args.write_phases, + keep_going=args.keep_going, + failure_log=args.out_root / "failures" / f"routing_failures{log_suffix}.jsonl", + sample_log=args.out_root / f"routing_sample_times{log_suffix}.jsonl", + resume_from=_routing_resume_sources(args.out_root, benchmark, args.model_key, args.layers, args), + ) + start = time.time() + summary = generate_true_routing_with_weights.run(routing_args) + summary["orchestrator_wall_seconds"] = time.time() - start + write_json(args.out_root / "routing_summaries" / f"p3_rev_{benchmark}_{args.model_key}.json", summary) + summaries.append(summary) + return summaries + + +def build_selected_traces(args: argparse.Namespace, selection: dict[str, Any]) -> list[Path]: + written: list[Path] = [] + for benchmark in args.benchmarks: + ids = _sample_ids_for_args(selection, benchmark, args) + if not ids: + continue + phases = "prefill,decode" if args.write_phases == "both" else args.write_phases + route_inputs = _routing_resume_sources(args.out_root, benchmark, args.model_key, args.layers, args) + if not route_inputs: + route_inputs = [_routing_output(args.out_root, benchmark, args.model_key, args.layers, args)] + for route_input in route_inputs: + build_args = argparse.Namespace( + input=route_input, + out_dir=args.out_root / "route_traces", + summary_out=args.out_root + / "route_trace_summaries" + / f"p3_rev_{benchmark}_{args.model_key}_{route_input.stem}.json", + sample_ids=",".join(ids), + layers=",".join(str(layer) for layer in args.layers), + phases=phases, + limit=None, + mlen=args.mlen, + blen=args.blen, + emu_threads=args.emu_threads, + allow_uniform_weights=False, + ) + written.extend(build_route_traces.build_traces(build_args)) + return written + + +def _trace_matches(trace: dict[str, Any], selection: dict[str, Any], layers: set[int], phases: set[str]) -> bool: + workload = trace.get("workload", {}) + benchmark = str(workload.get("benchmark")) + sample_id = str(workload.get("sample_id")) + selected = set(_sample_ids(selection, benchmark)) + if sample_id not in selected: + return False + if int(trace.get("model", {}).get("layer_index")) not in layers: + return False + if str(workload.get("phase")) not in phases: + return False + return True + + +def selected_trace_paths(args: argparse.Namespace, selection: dict[str, Any]) -> list[Path]: + layers = set(args.layers) + phases = {"decode"} if args.write_phases == "decode" else {"prefill"} if args.write_phases == "prefill" else {"prefill", "decode"} + paths = [] + for path in sorted((args.out_root / "route_traces").glob("*.json")): + try: + trace = load_json(path) + except Exception: + continue + if _trace_matches(trace, selection, layers, phases): + paths.append(path) + def key(path: Path) -> tuple[str, int, str, int]: + trace = load_json(path) + workload = trace["workload"] + return ( + str(workload.get("benchmark")), + int(workload.get("sample_index") or 0), + str(workload.get("sample_id")), + int(trace["model"]["layer_index"]), + ) + return sorted(paths, key=key) + + +def _filtered_trace_paths_by_id(trace_paths: list[Path], trace_id_file: Path | None) -> list[Path]: + if trace_id_file is None: + return trace_paths + wanted = { + line.strip() + for line in trace_id_file.read_text(encoding="utf-8").splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + if not wanted: + return [] + filtered = [] + for path in trace_paths: + try: + trace = load_json(path) + except Exception: + continue + if str(trace.get("trace_id")) in wanted: + filtered.append(path) + return filtered + + +def progress_snapshot(args: argparse.Namespace, selection: dict[str, Any]) -> dict[str, Any]: + selected_pairs = { + (benchmark, str(row["sample_id"])) + for benchmark, obj in selection.get("benchmarks", {}).items() + for row in obj.get("selected_samples", []) + } + route_progress = {} + for benchmark, obj in selection.get("benchmarks", {}).items(): + ids = {str(row["sample_id"]) for row in obj.get("selected_samples", [])} + route_file = _routing_output(args.out_root, benchmark, args.model_key, args.layers, args) + layer_text = "_".join(str(layer) for layer in args.layers) + route_files = sorted( + (args.out_root / "true_routing").glob( + f"{benchmark}_{args.model_key}_layers_{layer_text}_{args.write_phases}_d{args.decode_steps}*.jsonl" + ) + ) + got: dict[str, set[tuple[str, int, int | None]]] = {} + for current_route_file in route_files: + if not current_route_file.exists(): + continue + for row in iter_jsonl(current_route_file): + sample_id = str(row.get("sample_id")) + if sample_id not in ids: + continue + got.setdefault(sample_id, set()).add( + ( + str(row.get("phase")), + int(row.get("layer")), + int(row["decode_step"]) if row.get("decode_step") is not None else None, + ) + ) + complete = sum(1 for sample_id in ids if len(got.get(sample_id, set())) >= len(args.layers)) + route_progress[benchmark] = { + "selected_samples": len(ids), + "complete_route_samples": complete, + "route_rows": sum(len(value) for value in got.values()), + "routing_output": str(route_file), + "routing_outputs_scanned": [str(path) for path in route_files], + } + + trace_paths = selected_trace_paths(args, selection) + replay_complete = 0 + replay_rows = [] + for path in trace_paths: + trace = load_json(path) + build_dir = args.out_root / "trace_replay" / trace["trace_id"] + status = "complete" if _result_ok(build_dir, require_stage_profile=args.stage_profile) else "missing" + if status == "complete": + replay_complete += 1 + replay_rows.append({"trace_id": trace["trace_id"], "status": status}) + + exported_runs = 0 + selected_export = args.out_root / "p3_rev_emulation_timing.csv" + if selected_export.exists(): + with selected_export.open(newline="", encoding="utf-8") as handle: + exported_runs = sum(1 for _ in csv.DictReader(handle)) + failure_counts = {} + failure_dir = args.out_root / "failures" + if failure_dir.exists(): + for path in sorted(failure_dir.glob("*.jsonl")): + failure_counts[path.name] = sum(1 for _ in path.open(encoding="utf-8")) + + payload = { + "schema_version": 1, + "updated_at_unix": time.time(), + "selected_samples": sum(len(obj.get("selected_samples", [])) for obj in selection.get("benchmarks", {}).values()), + "expected_representative_runs": len(selected_pairs) * len(args.layers), + "route_progress": route_progress, + "selected_trace_files": len(trace_paths), + "replay_complete_runs": replay_complete, + "exported_selected_runs": exported_runs, + "failure_counts": failure_counts, + } + write_json(args.out_root / "p3_rev_progress_snapshot.json", payload) + return payload + + +def _result_ok(build_dir: Path, require_stage_profile: bool = True) -> bool: + result_path = build_dir / "qwen3_trace_replay_results.json" + run_stats_path = build_dir / "rust_emulator_run_stats.json" + if not result_path.exists() or not run_stats_path.exists(): + return False + if require_stage_profile and not (build_dir / "stage_profile.json").exists(): + return False + try: + result = load_json(result_path) + except Exception: + return False + return bool(result.get("functional_gate", True)) + + +def _run_trace_command( + trace_path: Path, + build_dir: Path, + args: argparse.Namespace, + *, + log_name: str, + use_time: bool = False, +) -> tuple[int, Path, int | None]: + cmd = [ + sys.executable, + "-m", + "transactional_emulator.testbench.window1_p2.qwen3_trace_replay_test", + str(trace_path), + "--build-dir", + str(build_dir), + "--mlen", + str(args.mlen), + "--blen", + str(args.blen), + "--emu-threads", + str(args.emu_threads), + ] + if args.stage_profile: + cmd.append("--stage-profile") + run_cmd = cmd + if use_time: + run_cmd = ["/usr/bin/time", "-v", *cmd] + build_dir.mkdir(parents=True, exist_ok=True) + log_path = build_dir / log_name + with log_path.open("w", encoding="utf-8") as log: + proc = subprocess.run( + run_cmd, + cwd=args.cwd, + text=True, + stdin=subprocess.DEVNULL, + stdout=log, + stderr=subprocess.STDOUT, + ) + rss_kb = None + if use_time: + text = log_path.read_text(encoding="utf-8", errors="replace") + match = re.search(r"Maximum resident set size \(kbytes\):\s*(\d+)", text) + if match: + rss_kb = int(match.group(1)) + return proc.returncode, log_path, rss_kb + + +def measure_rss(args: argparse.Namespace, trace_paths: list[Path]) -> dict[str, Any]: + if not trace_paths: + raise ValueError("no selected traces available for RSS probe") + probe_trace = trace_paths[0] + trace = load_json(probe_trace) + probe_root = args.out_root / "p3_rev_rss_probe" + build_dir = probe_root / trace["trace_id"] + start = time.time() + code, log_path, rss_kb = _run_trace_command(probe_trace, build_dir, args, log_name="rss_probe.log", use_time=True) + row = { + "schema_version": 1, + "trace": str(probe_trace), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "return_code": code, + "passed": code == 0, + "host_seconds": time.time() - start, + "max_rss_kb": rss_kb, + "log_path": str(log_path), + } + if code == 0 and args.prune_success_artifacts: + row.update(_prune_success_artifacts(build_dir)) + write_json(args.out_root / "p3_rev_rss_probe.json", row) + if code != 0: + raise RuntimeError(f"RSS probe failed for {trace['trace_id']}; see {log_path}") + return row + + +def _mem_available_kb() -> int: + for line in Path("/proc/meminfo").read_text(encoding="utf-8").splitlines(): + if line.startswith("MemAvailable:"): + return int(line.split()[1]) + return 0 + + +def choose_width(args: argparse.Namespace, rss: dict[str, Any] | None) -> dict[str, Any]: + if args.width is not None: + width = max(1, int(args.width)) + reason = "explicit_width" + else: + rss_kb = int((rss or {}).get("max_rss_kb") or 0) + cores = os.cpu_count() or 1 + mem_available_kb = _mem_available_kb() + memory_budget_kb = int(mem_available_kb * float(args.memory_budget_fraction)) + if rss_kb <= 0: + width = 1 + reason = "missing_rss_fallback" + else: + width = max(1, min(max(1, cores - 2), max(1, memory_budget_kb // rss_kb))) + reason = "min_cores_minus_2_and_memory_budget_over_rss" + if args.max_width is not None: + width = min(width, max(1, int(args.max_width))) + reason += "_capped" + payload = { + "schema_version": 1, + "width": width, + "reason": reason, + "cpu_count": os.cpu_count(), + "mem_available_kb": _mem_available_kb(), + "memory_budget_fraction": args.memory_budget_fraction, + "rss_probe": rss, + } + write_json(args.out_root / "p3_rev_parallel_width.json", payload) + return payload + + +def _worker_run_trace(task: dict[str, Any]) -> dict[str, Any]: + trace_path = Path(task["trace_path"]) + out_root = Path(task["out_root"]) + cwd = Path(task["cwd"]) + trace = load_json(trace_path) + build_dir = out_root / "trace_replay" / trace["trace_id"] + start = time.time() + row = { + "trace": str(trace_path), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "worker_pid": os.getpid(), + } + try: + if task["skip_existing"] and _result_ok(build_dir, require_stage_profile=task["stage_profile"]): + row.update({"status": "skipped_existing", "host_batch_seconds": 0.0}) + return row + ns = argparse.Namespace( + cwd=cwd, + mlen=task["mlen"], + blen=task["blen"], + emu_threads=task["emu_threads"], + stage_profile=task["stage_profile"], + ) + code, log_path, _ = _run_trace_command( + trace_path, + build_dir, + ns, + log_name="p3_rev_parallel_stdout.log", + use_time=False, + ) + row.update( + { + "return_code": code, + "status": "passed" if code == 0 else "failed", + "host_batch_seconds": time.time() - start, + "log_path": str(log_path), + } + ) + if row["status"] == "passed" and task["prune_success_artifacts"]: + row.update(_prune_success_artifacts(build_dir)) + except Exception as exc: + row.update( + { + "status": "failed", + "host_batch_seconds": time.time() - start, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + } + ) + return row + + +def _determinism_worker(task: dict[str, Any]) -> dict[str, Any]: + trace_path = Path(task["trace_path"]) + build_root = Path(task["build_root"]) + cwd = Path(task["cwd"]) + trace = load_json(trace_path) + build_dir = build_root / trace["trace_id"] + start = time.time() + row = { + "trace": str(trace_path), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "worker_pid": os.getpid(), + } + try: + if task.get("skip_existing") and _result_ok(build_dir, require_stage_profile=task["stage_profile"]): + row.update({"status": "skipped_existing", "host_batch_seconds": 0.0}) + if task.get("prune_success_artifacts"): + row.update(_prune_success_artifacts(build_dir)) + return row + ns = argparse.Namespace( + cwd=cwd, + mlen=task["mlen"], + blen=task["blen"], + emu_threads=task["emu_threads"], + stage_profile=task["stage_profile"], + ) + code, log_path, _ = _run_trace_command( + trace_path, + build_dir, + ns, + log_name=task["log_name"], + use_time=False, + ) + row.update( + { + "return_code": code, + "status": "passed" if code == 0 else "failed", + "host_batch_seconds": time.time() - start, + "log_path": str(log_path), + } + ) + if row["status"] == "passed" and task.get("prune_success_artifacts"): + row.update(_prune_success_artifacts(build_dir)) + except Exception as exc: + row.update( + { + "status": "failed", + "host_batch_seconds": time.time() - start, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + } + ) + return row + + +def _sha256_file(path: Path) -> str | None: + if not path.exists(): + return None + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +_VOLATILE_JSON_KEYS = { + "build_dir", + "command", + "config_path", + "cwd", + "ended_at_utc", + "host_wall_time_seconds", + "log_path", + "stage_profile_path", + "started_at_utc", + "stats_path", +} + + +def _normalized_json(value: Any) -> Any: + if isinstance(value, dict): + return { + key: _normalized_json(child) + for key, child in sorted(value.items()) + if key not in _VOLATILE_JSON_KEYS + } + if isinstance(value, list): + return [_normalized_json(child) for child in value] + return value + + +def _sha256_json(path: Path, *, normalized: bool) -> str | None: + if not path.exists(): + return None + payload = load_json(path) + if normalized: + payload = _normalized_json(payload) + text = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _comparison_hash(path: Path, mode: str) -> str | None: + if mode == "raw": + return _sha256_file(path) + if mode == "json": + return _sha256_json(path, normalized=False) + if mode == "normalized_json": + return _sha256_json(path, normalized=True) + raise ValueError(f"unknown comparison mode: {mode}") + + +def _determinism_trace_subset(trace_paths: list[Path], sample_count: int) -> list[Path]: + if sample_count <= 0: + return [] + by_sample: dict[tuple[str, str], list[Path]] = defaultdict(list) + order: list[tuple[str, str]] = [] + for path in trace_paths: + trace = load_json(path) + workload = trace["workload"] + key = (str(workload.get("benchmark")), str(workload.get("sample_id"))) + if key not in by_sample: + order.append(key) + by_sample[key].append(path) + chosen: list[Path] = [] + for key in order[:sample_count]: + chosen.extend(sorted(by_sample[key], key=lambda path: int(load_json(path)["model"]["layer_index"]))) + return chosen + + +def run_determinism_gate(args: argparse.Namespace, trace_paths: list[Path], width_info: dict[str, Any] | None) -> dict[str, Any]: + chosen = _determinism_trace_subset(trace_paths, args.determinism_samples) + if not chosen: + payload = { + "schema_version": 1, + "status": "skipped", + "reason": "no determinism traces selected", + "determinism_samples": args.determinism_samples, + } + write_json(args.out_root / "p3_rev_determinism.json", payload) + return payload + + root = args.out_root / "p3_rev_determinism" + serial_root = root / "serial" + parallel_root = root / "parallel" + compare_files = [ + {"file": "qwen3_trace_replay_results.json", "mode": "normalized_json"}, + {"file": "stage_profile.json", "mode": "json"}, + {"file": "rust_emulator_run_stats.json", "mode": "normalized_json"}, + {"file": "gather_scatter_results.json", "mode": "normalized_json"}, + ] + + serial_rows = [] + started = time.time() + for path in chosen: + row = _determinism_worker( + { + "trace_path": str(path), + "build_root": str(serial_root), + "cwd": str(args.cwd), + "mlen": args.mlen, + "blen": args.blen, + "emu_threads": args.emu_threads, + "stage_profile": args.stage_profile, + "log_name": "serial_stdout.log", + "skip_existing": True, + "prune_success_artifacts": args.prune_success_artifacts, + } + ) + serial_rows.append(row) + if row.get("status") != "passed" and not args.keep_going: + break + + parallel_width = args.determinism_width + if parallel_width is None: + parallel_width = min(int((width_info or {}).get("width") or 1), len(chosen), 4) + parallel_rows = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, parallel_width)) as executor: + futures = [ + executor.submit( + _determinism_worker, + { + "trace_path": str(path), + "build_root": str(parallel_root), + "cwd": str(args.cwd), + "mlen": args.mlen, + "blen": args.blen, + "emu_threads": args.emu_threads, + "stage_profile": args.stage_profile, + "log_name": "parallel_stdout.log", + "skip_existing": True, + "prune_success_artifacts": args.prune_success_artifacts, + }, + ) + for path in chosen + ] + for future in concurrent.futures.as_completed(futures): + parallel_rows.append(future.result()) + + comparisons = [] + for path in chosen: + trace = load_json(path) + trace_id = trace["trace_id"] + file_rows = [] + for item in compare_files: + name = item["file"] + mode = item["mode"] + serial_hash = _comparison_hash(serial_root / trace_id / name, mode) + parallel_hash = _comparison_hash(parallel_root / trace_id / name, mode) + file_rows.append( + { + "file": name, + "mode": mode, + "serial_sha256": serial_hash, + "parallel_sha256": parallel_hash, + "matches": serial_hash is not None and serial_hash == parallel_hash, + } + ) + comparisons.append( + { + "trace_id": trace_id, + "trace": str(path), + "files": file_rows, + "all_files_match": all(row["matches"] for row in file_rows), + } + ) + + payload = { + "schema_version": 1, + "status": "passed" if all(row["all_files_match"] for row in comparisons) else "failed", + "started_at": started, + "ended_at": time.time(), + "wall_seconds": time.time() - started, + "determinism_samples": args.determinism_samples, + "trace_count": len(chosen), + "parallel_width": parallel_width, + "serial_status_counts": dict(Counter(str(row.get("status")) for row in serial_rows)), + "parallel_status_counts": dict(Counter(str(row.get("status")) for row in parallel_rows)), + "compare_files": compare_files, + "comparisons": comparisons, + } + write_json(args.out_root / "p3_rev_determinism.json", payload) + return payload + + +def _progress_payload( + *, + args: argparse.Namespace, + run_id: str, + total: int, + started_at: float, + rows: list[dict[str, Any]], + running: int, + width: int, + manifest_jsonl: Path, + failure_log: Path, +) -> dict[str, Any]: + status_counts = Counter(str(row.get("status", "unknown")) for row in rows) + completed = sum(status_counts.values()) + return { + "schema_version": 1, + "run_id": run_id, + "pid": os.getpid(), + "updated_at_unix": time.time(), + "started_at_unix": started_at, + "elapsed_seconds": time.time() - started_at, + "selected_trace_count": total, + "completed_events": completed, + "running": running, + "remaining_events": max(0, total - completed - running), + "status_counts": dict(status_counts), + "width": width, + "manifest_jsonl": str(manifest_jsonl), + "failure_log": str(failure_log), + } + + +def _replay_artifact_paths(args: argparse.Namespace) -> dict[str, Path]: + name = getattr(args, "replay_run_name", None) or "p3_rev_parallel_replay" + if name == "p3_rev_parallel_replay": + return { + "progress": args.out_root / "p3_rev_progress.json", + "manifest_jsonl": args.out_root / "p3_rev_parallel_replay_manifest.jsonl", + "manifest_json": args.out_root / "p3_rev_parallel_replay_manifest.json", + "failure_log": args.out_root / "failures" / "p3_rev_replay_failures.jsonl", + } + safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", name).strip("_") or "p3_rev_parallel_replay" + return { + "progress": args.out_root / f"{safe}_progress.json", + "manifest_jsonl": args.out_root / f"{safe}_manifest.jsonl", + "manifest_json": args.out_root / f"{safe}_manifest.json", + "failure_log": args.out_root / "failures" / f"{safe}_failures.jsonl", + } + + +def run_parallel_replay(args: argparse.Namespace, trace_paths: list[Path], width_info: dict[str, Any]) -> dict[str, Any]: + run_id = time.strftime("%Y%m%d_%H%M%S") + width = int(width_info["width"]) + started_at = time.time() + rows: list[dict[str, Any]] = [] + paths = _replay_artifact_paths(args) + manifest_jsonl = paths["manifest_jsonl"] + failure_log = paths["failure_log"] + args.out_root.mkdir(parents=True, exist_ok=True) + (args.out_root / "failures").mkdir(parents=True, exist_ok=True) + tasks = [ + { + "trace_path": str(path), + "out_root": str(args.out_root), + "cwd": str(args.cwd), + "mlen": args.mlen, + "blen": args.blen, + "emu_threads": args.emu_threads, + "stage_profile": args.stage_profile, + "skip_existing": args.skip_existing, + "prune_success_artifacts": args.prune_success_artifacts, + } + for path in trace_paths + ] + progress_path = paths["progress"] + _atomic_json( + progress_path, + _progress_payload( + args=args, + run_id=run_id, + total=len(tasks), + started_at=started_at, + rows=rows, + running=min(width, len(tasks)), + width=width, + manifest_jsonl=manifest_jsonl, + failure_log=failure_log, + ), + ) + with concurrent.futures.ThreadPoolExecutor(max_workers=width) as executor: + futures = [executor.submit(_worker_run_trace, task) for task in tasks] + for future in concurrent.futures.as_completed(futures): + row = future.result() + row["run_id"] = run_id + rows.append(row) + _append_jsonl(manifest_jsonl, row) + if row.get("status") == "failed": + _append_jsonl(failure_log, row) + running = sum(1 for item in futures if item.running()) + _atomic_json( + progress_path, + _progress_payload( + args=args, + run_id=run_id, + total=len(tasks), + started_at=started_at, + rows=rows, + running=running, + width=width, + manifest_jsonl=manifest_jsonl, + failure_log=failure_log, + ), + ) + manifest = { + "schema_version": 1, + "run_id": run_id, + "started_at": started_at, + "ended_at": time.time(), + "orchestrator_wall_seconds": time.time() - started_at, + "trace_count": len(trace_paths), + "width": width, + "status_counts": dict(Counter(str(row.get("status", "unknown")) for row in rows)), + "runs": rows, + } + write_json(paths["manifest_json"], manifest) + return manifest + + +def _sum_stages(profile: dict[str, Any], names: set[str]) -> int: + stages = profile.get("stages", {}) + if not isinstance(stages, dict): + return 0 + return sum(int(stages.get(name, {}).get("wall_cycles") or 0) for name in names) + + +def export_selected(args: argparse.Namespace, selection: dict[str, Any]) -> dict[str, Any]: + selected: set[tuple[str, str]] = set() + for benchmark in args.benchmarks: + selected.update((benchmark, sample_id) for sample_id in _sample_ids(selection, benchmark)) + export_args = argparse.Namespace( + root=args.out_root / "trace_replay", + out_prefix=args.out_root / "_p3_rev_unfiltered", + out_json=args.out_root / "_p3_rev_unfiltered_timing_summary.json", + ) + payload = export_pilot_results.export(export_args) + runs = [row for row in payload["runs"] if (row["benchmark"], str(row["sample_id"])) in selected] + trace_ids = {row["trace_id"] for row in runs} + stage_stack = [row for row in payload["stage_stack"] if row["trace_id"] in trace_ids] + routing_tax = [row for row in payload["routing_tax"] if row["trace_id"] in trace_ids] + bytes_validation = [row for row in payload["bytes_validation"] if row["trace_id"] in trace_ids] + summary_rows = _summary_rows(runs, routing_tax, bytes_validation) + prefix = args.out_root / "p3_rev_emulation" + write_csv(prefix.with_name(prefix.name + "_timing.csv"), runs) + write_csv(prefix.with_name(prefix.name + "_stage_stack.csv"), stage_stack) + write_csv(prefix.with_name(prefix.name + "_routing_tax.csv"), routing_tax) + write_csv(prefix.with_name(prefix.name + "_bytes_validation.csv"), bytes_validation) + write_csv(prefix.with_name(prefix.name + "_distribution_stats.csv"), summary_rows) + filtered = { + "schema_version": 1, + "root": str(args.out_root / "trace_replay"), + "selected_samples": str(args.selected_out), + "measurement_note": MEASUREMENT_NOTE, + "runs": runs, + "stage_stack": stage_stack, + "routing_tax": routing_tax, + "bytes_validation": bytes_validation, + "distribution_stats": summary_rows, + "unfiltered_runs": len(payload["runs"]), + } + write_json(args.out_root / "p3_rev_emulation_timing_summary.json", filtered) + return filtered + + +def _summary_rows( + runs: list[dict[str, Any]], + routing_tax: list[dict[str, Any]], + bytes_validation: list[dict[str, Any]], +) -> list[dict[str, Any]]: + tax_by_trace = {row["trace_id"]: row for row in routing_tax} + bytes_by_trace = {row["trace_id"]: row for row in bytes_validation} + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + for row in runs: + benchmark = str(row.get("benchmark", "")) + category = str(row.get("category") or "uncategorized") + groups[(benchmark, "__all__")].append(row) + groups[(benchmark, category)].append(row) + out = [] + for (benchmark, category), rows in sorted(groups.items()): + trace_ids = [row["trace_id"] for row in rows] + cycle_stats = _stats([row.get("sim_latency_cycles") for row in rows]) + tax_stats = _stats([tax_by_trace.get(trace_id, {}).get("routing_tax_cycles") for trace_id in trace_ids]) + read_stats = _stats( + [bytes_by_trace.get(trace_id, {}).get("physical_hbm_bytes_read_run_stats") for trace_id in trace_ids] + ) + write_stats = _stats( + [bytes_by_trace.get(trace_id, {}).get("physical_hbm_bytes_written_run_stats") for trace_id in trace_ids] + ) + out.append( + { + "benchmark": benchmark, + "category": category, + "runs": len(rows), + "cycles_min": cycle_stats["min"], + "cycles_median": cycle_stats["median"], + "cycles_mean": cycle_stats["mean"], + "cycles_p95": cycle_stats["p95"], + "cycles_max": cycle_stats["max"], + "routing_tax_cycles_min": tax_stats["min"], + "routing_tax_cycles_median": tax_stats["median"], + "routing_tax_cycles_mean": tax_stats["mean"], + "routing_tax_cycles_p95": tax_stats["p95"], + "routing_tax_cycles_max": tax_stats["max"], + "hbm_read_bytes_min": read_stats["min"], + "hbm_read_bytes_median": read_stats["median"], + "hbm_read_bytes_mean": read_stats["mean"], + "hbm_read_bytes_p95": read_stats["p95"], + "hbm_read_bytes_max": read_stats["max"], + "hbm_written_bytes_min": write_stats["min"], + "hbm_written_bytes_median": write_stats["median"], + "hbm_written_bytes_mean": write_stats["mean"], + "hbm_written_bytes_p95": write_stats["p95"], + "hbm_written_bytes_max": write_stats["max"], + "measurement_note": MEASUREMENT_NOTE, + } + ) + return out + + +def _line_count(path: Path) -> int | None: + if not path.exists(): + return None + with path.open("r", encoding="utf-8", errors="ignore") as fh: + return sum(1 for _ in fh) + + +def _failure_counts(args: argparse.Namespace) -> dict[str, Any]: + paths: dict[str, Path] = {} + for path in sorted((args.out_root / "failures").glob("routing_failures*.jsonl")): + paths[path.name] = path + for path in [ + args.out_root / "failures" / "replay_failures.jsonl", + args.out_root / "failures" / "p3_rev_replay_failures.jsonl", + _replay_artifact_paths(args)["failure_log"], + ]: + paths[path.name] = path + counts = {name: _line_count(path) for name, path in sorted(paths.items())} + return { + "total_existing_failures": sum(count for count in counts.values() if count is not None), + "by_file": counts, + } + + +def _report_artifacts(args: argparse.Namespace) -> dict[str, str]: + prefix = args.out_root / "p3_rev_emulation" + replay_paths = _replay_artifact_paths(args) + return { + "selected_samples": str(args.selected_out), + "timing_csv": str(prefix.with_name(prefix.name + "_timing.csv")), + "stage_stack_csv": str(prefix.with_name(prefix.name + "_stage_stack.csv")), + "routing_tax_csv": str(prefix.with_name(prefix.name + "_routing_tax.csv")), + "bytes_validation_csv": str(prefix.with_name(prefix.name + "_bytes_validation.csv")), + "distribution_stats_csv": str(prefix.with_name(prefix.name + "_distribution_stats.csv")), + "timing_summary_json": str(args.out_root / "p3_rev_emulation_timing_summary.json"), + "replay_progress_json": str(replay_paths["progress"]), + "replay_manifest_jsonl": str(replay_paths["manifest_jsonl"]), + "replay_manifest_json": str(replay_paths["manifest_json"]), + "replay_failure_log": str(replay_paths["failure_log"]), + "report_json": str(args.out_root / f"{args.report_name}.json"), + "report_md": str(args.out_root / f"{args.report_name}.md"), + } + + +def _progress_summary(progress: dict[str, Any] | None) -> dict[str, Any] | None: + if not progress: + return None + route_progress = {} + for benchmark, row in (progress.get("route_progress") or {}).items(): + route_progress[benchmark] = { + "selected_samples": row.get("selected_samples"), + "complete_route_samples": row.get("complete_route_samples"), + "route_rows": row.get("route_rows"), + } + return { + "selected_samples": progress.get("selected_samples"), + "expected_representative_runs": progress.get("expected_representative_runs"), + "selected_trace_files": progress.get("selected_trace_files"), + "replay_complete_runs": progress.get("replay_complete_runs"), + "exported_selected_runs": progress.get("exported_selected_runs"), + "failure_counts": progress.get("failure_counts"), + "route_progress": route_progress, + } + + +def write_report( + args: argparse.Namespace, + selection: dict[str, Any], + trace_paths: list[Path], + rss: dict[str, Any] | None, + width: dict[str, Any] | None, + replay_manifest: dict[str, Any] | None, + export_payload: dict[str, Any] | None, + determinism: dict[str, Any] | None = None, + progress: dict[str, Any] | None = None, +) -> dict[str, Any]: + selected_counts = { + benchmark: len(_sample_ids(selection, benchmark)) + for benchmark in args.benchmarks + } + expected_runs = sum(selected_counts.values()) * len(args.layers) + exported_runs = len((export_payload or {}).get("runs", [])) + functional_counts = Counter(str(row.get("functional_gate")) for row in (export_payload or {}).get("runs", [])) + cycle_counts = Counter(str(row.get("cycle_accounting_status")) for row in (export_payload or {}).get("runs", [])) + byte_rows = (export_payload or {}).get("bytes_validation", []) + read_match = Counter(str(row.get("physical_read_matches")) for row in byte_rows) + write_match = Counter(str(row.get("physical_written_matches")) for row in byte_rows) + replay_counts = dict(Counter(str(row.get("status", "unknown")) for row in (replay_manifest or {}).get("runs", []))) + replay_progress = None + progress_path = _replay_artifact_paths(args)["progress"] + if progress_path.exists(): + try: + replay_progress = load_json(progress_path) + except Exception: + replay_progress = None + if not replay_counts and replay_progress: + replay_counts = dict(replay_progress.get("status_counts") or {}) + report = { + "schema_version": 1, + "generated_at_unix": time.time(), + "selected_samples": str(args.selected_out), + "selected_counts": selected_counts, + "expected_representative_runs": expected_runs, + "route_trace_count": len(trace_paths), + "exported_runs": exported_runs, + "rss_probe": rss, + "parallel_width": width, + "determinism": determinism, + "progress_snapshot": progress, + "progress_summary": _progress_summary(progress), + "replay_progress": replay_progress, + "replay_status_counts_latest_run": replay_counts, + "functional_gate_counts": dict(functional_counts), + "cycle_accounting_counts": dict(cycle_counts), + "physical_read_match_counts": dict(read_match), + "physical_written_match_counts": dict(write_match), + "failure_counts": _failure_counts(args), + "artifacts": _report_artifacts(args), + "measurement_note": MEASUREMENT_NOTE, + "known_limits": [ + "P3-rev reuses P2 fixed-route replay semantics.", + "Numbers are self-consistent upper bounds pending RTL calibration.", + "Replay is decode tok1 only.", + "Route replay uses zero-activation timing shape, as in P2.", + "Routing tax is stage-derived and excludes device router GEMM/top-k.", + ], + } + write_json(args.out_root / f"{args.report_name}.json", report) + _write_report_md(args.out_root / f"{args.report_name}.md", report, selection) + return report + + +def _write_report_md(path: Path, report: dict[str, Any], selection: dict[str, Any]) -> None: + width = report.get("parallel_width") or {} + rss = report.get("rss_probe") or {} + lines = [ + "# Window 1 P3-rev Representative Timing Report", + "", + f"- measurement note: {report['measurement_note']}", + f"- selected samples: `{report['selected_samples']}`", + f"- selected counts: `{report['selected_counts']}`", + f"- expected representative runs: `{report['expected_representative_runs']}`", + f"- route trace files currently available: `{report['route_trace_count']}`", + f"- exported selected runs: `{report['exported_runs']}`", + f"- progress summary: `{report.get('progress_summary')}`", + "", + "## Subset", + "", + ] + for benchmark, bench in selection.get("benchmarks", {}).items(): + lines.append( + f"- {benchmark}: selected `{bench.get('selected_count')}` / available `{bench.get('available')}`, " + f"method `{bench.get('selection_method')}`" + ) + lines.extend( + [ + "", + "## Parallel Replay", + "", + f"- RSS probe trace: `{rss.get('trace_id')}`", + f"- RSS probe max RSS: `{_fmt_bytes(int(rss.get('max_rss_kb') or 0) * 1024 if rss else None)}`", + f"- RSS probe wall time: `{_fmt_seconds(rss.get('host_seconds'))}`", + f"- selected width: `{width.get('width')}`", + f"- width reason: `{width.get('reason')}`", + f"- latest replay status counts: `{report['replay_status_counts_latest_run']}`", + f"- replay completed/running/remaining: " + f"`{(report.get('replay_progress') or {}).get('completed_events')}` / " + f"`{(report.get('replay_progress') or {}).get('running')}` / " + f"`{(report.get('replay_progress') or {}).get('remaining_events')}`", + f"- replay progress file: `{(report.get('artifacts') or {}).get('replay_progress_json')}`", + f"- determinism status: `{(report.get('determinism') or {}).get('status')}`", + f"- determinism trace count: `{(report.get('determinism') or {}).get('trace_count')}`", + "", + "## Gates", + "", + f"- functional gate counts: `{report['functional_gate_counts']}`", + f"- cycle accounting counts: `{report['cycle_accounting_counts']}`", + f"- physical read match counts: `{report['physical_read_match_counts']}`", + f"- physical written match counts: `{report['physical_written_match_counts']}`", + f"- failure counts: `{report['failure_counts']}`", + "", + "## Artifacts", + "", + f"- timing CSV: `{(report.get('artifacts') or {}).get('timing_csv')}`", + f"- stage stack CSV: `{(report.get('artifacts') or {}).get('stage_stack_csv')}`", + f"- routing tax CSV: `{(report.get('artifacts') or {}).get('routing_tax_csv')}`", + f"- bytes validation CSV: `{(report.get('artifacts') or {}).get('bytes_validation_csv')}`", + f"- distribution stats CSV: `{(report.get('artifacts') or {}).get('distribution_stats_csv')}`", + "", + "## Limits", + "", + ] + ) + lines.extend(f"- {item}" for item in report["known_limits"]) + lines.append("") + path.write_text("\n".join(lines), encoding="utf-8") + + +def _write_stage_status(args: argparse.Namespace, stage: str, status: str, **extra: Any) -> None: + if getattr(args, "no_stage_status", False): + return + payload = { + "schema_version": 1, + "pid": os.getpid(), + "updated_at_unix": time.time(), + "stage": stage, + "status": status, + **extra, + } + _atomic_json(args.out_root / "p3_rev_stage_status.json", payload) + + +def run(args: argparse.Namespace) -> dict[str, Any]: + ensure_paths() + args.out_root.mkdir(parents=True, exist_ok=True) + (args.out_root / "failures").mkdir(parents=True, exist_ok=True) + _write_stage_status(args, "start", "running", actions=sorted(args.actions)) + selection = _load_selection(args) + _write_stage_status( + args, + "select", + "completed", + selected_counts={benchmark: len(_sample_ids(selection, benchmark)) for benchmark in args.benchmarks}, + ) + routing_summaries = None + if "route" in args.actions: + _write_stage_status(args, "route", "running") + routing_summaries = run_route_generation(args, selection) + write_json(args.out_root / "p3_rev_routing_summaries.json", {"summaries": routing_summaries}) + _write_stage_status(args, "route", "completed", summaries=routing_summaries) + if "build" in args.actions: + _write_stage_status(args, "build", "running") + build_selected_traces(args, selection) + _write_stage_status(args, "build", "completed") + progress = progress_snapshot(args, selection) + if "progress" in args.actions: + _write_stage_status(args, "progress", "completed", progress=progress) + trace_paths = selected_trace_paths(args, selection) + replay_trace_paths = _filtered_trace_paths_by_id(trace_paths, args.replay_trace_id_file) + rss = None + if "rss" in args.actions: + _write_stage_status(args, "rss", "running", selected_trace_count=len(trace_paths)) + rss = measure_rss(args, trace_paths) + _write_stage_status(args, "rss", "completed", rss_probe=rss) + elif (args.out_root / "p3_rev_rss_probe.json").exists(): + rss = load_json(args.out_root / "p3_rev_rss_probe.json") + width = None + if "width" in args.actions or "replay" in args.actions: + _write_stage_status(args, "width", "running") + width = choose_width(args, rss) + _write_stage_status(args, "width", "completed", width=width) + elif (args.out_root / "p3_rev_parallel_width.json").exists(): + width = load_json(args.out_root / "p3_rev_parallel_width.json") + determinism = None + if "determinism" in args.actions: + if width is None: + width = choose_width(args, rss) + _write_stage_status( + args, + "determinism", + "running", + selected_trace_count=len(trace_paths), + determinism_samples=args.determinism_samples, + ) + determinism = run_determinism_gate(args, trace_paths, width) + _write_stage_status(args, "determinism", "completed", determinism=determinism) + elif (args.out_root / "p3_rev_determinism.json").exists(): + determinism = load_json(args.out_root / "p3_rev_determinism.json") + replay_manifest = None + if "replay" in args.actions: + if width is None: + width = choose_width(args, rss) + _write_stage_status(args, "replay", "running", selected_trace_count=len(replay_trace_paths), width=width) + replay_manifest = run_parallel_replay(args, replay_trace_paths, width) + _write_stage_status(args, "replay", "completed", replay_manifest=replay_manifest) + elif _replay_artifact_paths(args)["manifest_json"].exists(): + replay_manifest = load_json(_replay_artifact_paths(args)["manifest_json"]) + elif args.replay_run_name == "p3_rev_parallel_replay" and (args.out_root / "p3_rev_parallel_replay_manifest.json").exists(): + replay_manifest = load_json(args.out_root / "p3_rev_parallel_replay_manifest.json") + export_payload = None + if "export" in args.actions: + _write_stage_status(args, "export", "running") + export_payload = export_selected(args, selection) + _write_stage_status(args, "export", "completed", exported_runs=len(export_payload.get("runs", []))) + elif (args.out_root / "p3_rev_emulation_timing_summary.json").exists(): + export_payload = load_json(args.out_root / "p3_rev_emulation_timing_summary.json") + _write_stage_status(args, "report", "running") + progress = progress_snapshot(args, selection) + report = write_report(args, selection, trace_paths, rss, width, replay_manifest, export_payload, determinism, progress) + _write_stage_status(args, "done", "completed", report=str(args.out_root / f"{args.report_name}.md")) + return report + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT) + parser.add_argument("--selected-out", type=Path) + parser.add_argument("--model-key", choices=sorted(MODEL_CONFIGS), default="qwen3") + parser.add_argument("--benchmarks", type=_parse_csv_strings, default=_parse_csv_strings("bfcl_v3,gpqa_diamond")) + parser.add_argument("--bfcl-target", type=int, default=300) + parser.add_argument("--seed", type=int, default=DEFAULT_SEED) + parser.add_argument("--reselect", action="store_true") + parser.add_argument("--layers", type=_parse_csv_ints, default=DEFAULT_LAYERS) + parser.add_argument("--write-phases", choices=("prefill", "decode", "both"), default="decode") + parser.add_argument("--decode-steps", type=int, default=1) + parser.add_argument( + "--actions", + type=lambda value: set(_parse_csv_strings(value)), + default={"select", "route", "build", "rss", "width", "determinism", "replay", "export", "report"}, + help="Comma-separated actions: select,route,build,progress,rss,width,determinism,replay,export,report", + ) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--max-input-tokens", type=int) + parser.add_argument("--sort-by-length", action="store_true") + parser.add_argument("--mlen", type=int, default=128) + parser.add_argument("--blen", type=int, default=4) + parser.add_argument("--emu-threads", type=int, default=1) + parser.add_argument("--cwd", type=Path, default=REPO_ROOT) + parser.add_argument("--stage-profile", action="store_true") + parser.add_argument("--skip-existing", action="store_true") + parser.add_argument("--keep-going", action="store_true") + parser.add_argument("--prune-success-artifacts", action="store_true") + parser.add_argument("--width", type=int) + parser.add_argument("--max-width", type=int) + parser.add_argument("--memory-budget-fraction", type=float, default=0.70) + parser.add_argument("--determinism-samples", type=int, default=20) + parser.add_argument("--determinism-width", type=int) + parser.add_argument("--no-stage-status", action="store_true") + parser.add_argument("--routing-log-suffix", default="") + parser.add_argument("--routing-output-suffix", default="") + parser.add_argument("--sample-shard-count", type=int) + parser.add_argument("--sample-shard-index", type=int) + parser.add_argument( + "--replay-trace-id-file", + type=Path, + help="Replay only trace ids listed one per line; defaults to all selected traces.", + ) + parser.add_argument( + "--replay-run-name", + default="p3_rev_parallel_replay", + help="Artifact namespace for replay progress/manifest files.", + ) + parser.add_argument("--report-name", default="p3_rev_report") + args = parser.parse_args() + if args.selected_out is None: + args.selected_out = args.out_root / "p3_rev_selected_samples.json" + args.layers = args.layers or DEFAULT_LAYERS + unknown_actions = set(args.actions) - { + "select", + "route", + "build", + "progress", + "rss", + "width", + "determinism", + "replay", + "export", + "report", + } + if unknown_actions: + raise ValueError(f"unknown actions: {sorted(unknown_actions)}") + unknown_benchmarks = [benchmark for benchmark in args.benchmarks if benchmark not in INPUT_FILES] + if unknown_benchmarks: + raise ValueError(f"unknown benchmarks: {unknown_benchmarks}") + report = run(args) + print(f"wrote {args.out_root / (args.report_name + '.md')}") + print( + "selected=" + f"{report['selected_counts']} " + f"traces={report['route_trace_count']} " + f"exported={report['exported_runs']}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p3/run_p3_scale_batch.py b/transactional_emulator/testbench/window1_p3/run_p3_scale_batch.py new file mode 100644 index 00000000..848fdba1 --- /dev/null +++ b/transactional_emulator/testbench/window1_p3/run_p3_scale_batch.py @@ -0,0 +1,722 @@ +#!/usr/bin/env python3 +"""Window 1 P3 scale runner built on top of the P2 route replay sidecars. + +This wrapper intentionally does not change emulator timing semantics. It +orchestrates the existing P2 steps: + +1. generate true router traces from local model weights, +2. convert true-routing rows to replay traces, +3. replay traces through the Rust emulator harness, +4. export timing/tax/bytes/stage tables, +5. write scale/cost/failure reports for P3. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import PLENA_ROOT +from transactional_emulator.testbench.window1_p2 import ( + build_route_traces, + export_pilot_results, + generate_true_routing_with_weights, + run_trace_batch, +) +from transactional_emulator.testbench.window1_p2.p2_utils import ( + INPUT_FILES, + MODEL_CONFIGS, + ensure_paths, + iter_jsonl, + load_json, + write_csv, + write_json, +) + + +DEFAULT_OUT_ROOT = PLENA_ROOT / "outputs" / "window1_p3" +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_LAYERS = "0,12,23" +FULL_SAMPLE_COUNTS = { + "bfcl_v3": 5251, + "gpqa_diamond": 198, +} +MEASUREMENT_NOTE = "self-consistent upper bound, absolute accuracy pending RTL (Window 2)" + + +def _parse_csv_ints(value: str) -> list[int]: + return [int(part) for part in value.split(",") if part.strip()] + + +def _parse_csv_strings(value: str) -> list[str]: + return [part.strip() for part in value.split(",") if part.strip()] + + +def _limit_for(args: argparse.Namespace, benchmark: str) -> int | None: + if benchmark == "bfcl_v3": + return args.bfcl_limit + if benchmark == "gpqa_diamond": + return args.gpqa_limit + raise ValueError(f"unknown benchmark: {benchmark}") + + +def _routing_output(out_root: Path, benchmark: str, model_key: str, layers: list[int], args: argparse.Namespace) -> Path: + layer_text = "_".join(str(layer) for layer in layers) + return ( + out_root + / "true_routing" + / f"{benchmark}_{model_key}_layers_{layer_text}_{args.write_phases}_d{args.decode_steps}.jsonl" + ) + + +def _route_trace_summary(out_root: Path, benchmark: str, model_key: str) -> Path: + return out_root / "route_trace_summaries" / f"{benchmark}_{model_key}_route_traces_summary.json" + + +def _routing_summary_path(out_root: Path, benchmark: str, model_key: str) -> Path: + return out_root / "routing_summaries" / f"{benchmark}_{model_key}_true_routing_summary.json" + + +def _routing_sample_log_path(out_root: Path) -> Path: + return out_root / "routing_sample_times.jsonl" + + +def _dir_size_bytes(path: Path) -> int: + if not path.exists(): + return 0 + if path.is_file(): + return path.stat().st_size + total = 0 + for child in path.rglob("*"): + if child.is_file(): + total += child.stat().st_size + return total + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _append_jsonl(path: Path, row: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _percentile(values: list[float], pct: float) -> float | None: + if not values: + return None + ordered = sorted(values) + if len(ordered) == 1: + return float(ordered[0]) + rank = (len(ordered) - 1) * pct + lo = int(rank) + hi = min(lo + 1, len(ordered) - 1) + frac = rank - lo + return float(ordered[lo] * (1.0 - frac) + ordered[hi] * frac) + + +def _stats(values: list[int | float | None]) -> dict[str, Any]: + clean = [float(value) for value in values if value is not None] + if not clean: + return { + "count": 0, + "min": None, + "median": None, + "mean": None, + "p95": None, + "max": None, + } + return { + "count": len(clean), + "min": min(clean), + "median": float(statistics.median(clean)), + "mean": float(statistics.mean(clean)), + "p95": _percentile(clean, 0.95), + "max": max(clean), + } + + +def _input_count(benchmark: str, model_key: str) -> int: + return sum(1 for row in iter_jsonl(INPUT_FILES[benchmark]) if row.get("model_key") == model_key) + + +def _completed_routing_samples(out_root: Path, args: argparse.Namespace, layers: list[int]) -> dict[str, int]: + counts: dict[str, int] = {} + for benchmark in args.benchmarks: + path = _routing_output(out_root, benchmark, args.model_key, layers, args) + sample_ids = {str(row.get("sample_id")) for row in iter_jsonl(path)} if path.exists() else set() + counts[benchmark] = len(sample_ids) + return counts + + +def _run_true_routing(args: argparse.Namespace, benchmark: str, layers: list[int], out_root: Path) -> dict[str, Any]: + output = _routing_output(out_root, benchmark, args.model_key, layers, args) + failure_log = out_root / "failures" / "routing_failures.jsonl" + sample_log = _routing_sample_log_path(out_root) + routing_args = argparse.Namespace( + benchmark=benchmark, + model_key=args.model_key, + limit=_limit_for(args, benchmark), + sample_ids=None, + max_input_tokens=args.max_input_tokens, + layers=layers, + decode_steps=args.decode_steps, + threads=args.threads, + batch_size=args.batch_size, + output=output, + resume=args.resume, + sort_by_length=args.sort_by_length, + write_phases=args.write_phases, + keep_going=args.keep_going, + failure_log=failure_log, + sample_log=sample_log, + ) + start = time.time() + summary = generate_true_routing_with_weights.run(routing_args) + summary["orchestrator_wall_seconds"] = time.time() - start + summary["input_total_samples_for_model"] = _input_count(benchmark, args.model_key) + write_json(_routing_summary_path(out_root, benchmark, args.model_key), summary) + return summary + + +def _build_route_traces(args: argparse.Namespace, benchmark: str, layers: list[int], out_root: Path) -> list[Path]: + input_path = _routing_output(out_root, benchmark, args.model_key, layers, args) + phases = "prefill,decode" if args.write_phases == "both" else args.write_phases + build_args = argparse.Namespace( + input=input_path, + out_dir=out_root / "route_traces", + summary_out=_route_trace_summary(out_root, benchmark, args.model_key), + sample_ids=None, + layers=",".join(str(layer) for layer in layers), + phases=phases, + limit=None, + mlen=args.mlen, + blen=args.blen, + emu_threads=args.emu_threads, + allow_uniform_weights=False, + ) + return build_route_traces.build_traces(build_args) + + +def _run_replay(args: argparse.Namespace, out_root: Path) -> dict[str, Any]: + replay_args = argparse.Namespace( + trace_glob=[str(out_root / "route_traces" / "*.json")], + root=out_root / "trace_replay", + manifest_out=out_root / "trace_replay_batch_manifest.json", + cwd=args.cwd, + mlen=args.mlen, + blen=args.blen, + emu_threads=args.emu_threads, + limit=args.replay_limit, + stage_profile=args.stage_profile, + skip_existing=args.skip_existing, + keep_going=args.keep_going, + keep_dumps=args.keep_dumps, + prune_success_artifacts=args.prune_heavy_artifacts, + experimental_overlap_prefetch_compute=False, + ) + start = time.time() + manifest = run_trace_batch.run(replay_args) + manifest["orchestrator_wall_seconds"] = time.time() - start + write_json(replay_args.manifest_out, manifest) + failures = [row for row in manifest.get("runs", []) if row.get("status") == "failed"] + _write_jsonl(out_root / "failures" / "replay_failures.jsonl", failures) + return manifest + + +def _export(args: argparse.Namespace, out_root: Path) -> dict[str, Any]: + prefix = out_root / "gpqa_bfcl_emulation" + export_args = argparse.Namespace( + root=out_root / "trace_replay", + out_prefix=prefix, + out_json=out_root / "gpqa_bfcl_emulation_timing_summary.json", + ) + payload = export_pilot_results.export(export_args) + _write_distribution_stats(out_root, payload) + return payload + + +def _prune_heavy_artifacts(out_root: Path) -> dict[str, Any]: + """Drop replay intermediates that are not needed by the P3 exported tables.""" + keep_names = { + "comparison_params.json", + "compile_info.json", + "generated_asm_code.asm", + "generated_machine_code.mem", + "gather_scatter_results.json", + "p2_batch_stdout.log", + "qwen3_trace_replay_manifest.json", + "qwen3_trace_replay_results.json", + "rust_emulator_run_stats.json", + "rust_emulator_stdout.log", + "stage_profile.json", + "tensor_layouts.json", + "trace.json", + } + prune_suffixes = {".pt"} + prune_names = { + "hbm_for_behave_sim.bin", + "fp_sram.bin", + "int_sram.bin", + "vram_preload.bin", + "vector_result.mem", + } + rows: list[dict[str, Any]] = [] + trace_root = out_root / "trace_replay" + for build_dir in sorted(path for path in trace_root.glob("*") if path.is_dir()): + result_path = build_dir / "qwen3_trace_replay_results.json" + if not result_path.exists(): + continue + for path in sorted(child for child in build_dir.iterdir() if child.is_file()): + if path.name in keep_names: + continue + if path.suffix not in prune_suffixes and path.name not in prune_names: + continue + size = path.stat().st_size + path.unlink() + rows.append({"build_dir": str(build_dir), "path": str(path), "bytes": size}) + total = sum(int(row["bytes"]) for row in rows) + summary = { + "schema_version": 1, + "pruned_file_count": len(rows), + "pruned_bytes": total, + "note": "Pruned only post-run intermediates; timing/result/stage/profile/trace/log evidence files are retained.", + "files": rows, + } + write_json(out_root / "pruned_artifacts_manifest.json", summary) + return summary + + +def _write_distribution_stats(out_root: Path, payload: dict[str, Any]) -> None: + tax_by_trace = {row["trace_id"]: row for row in payload.get("routing_tax", [])} + bytes_by_trace = {row["trace_id"]: row for row in payload.get("bytes_validation", [])} + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + for row in payload.get("runs", []): + benchmark = str(row.get("benchmark", "")) + category = str(row.get("category") or "uncategorized") + groups[(benchmark, "__all__")].append(row) + groups[(benchmark, category)].append(row) + + stat_rows: list[dict[str, Any]] = [] + for (benchmark, category), rows in sorted(groups.items()): + trace_ids = [row["trace_id"] for row in rows] + cycle_stats = _stats([row.get("sim_latency_cycles") for row in rows]) + tax_stats = _stats([tax_by_trace.get(trace_id, {}).get("routing_tax_cycles") for trace_id in trace_ids]) + read_stats = _stats( + [bytes_by_trace.get(trace_id, {}).get("physical_hbm_bytes_read_run_stats") for trace_id in trace_ids] + ) + write_stats = _stats( + [bytes_by_trace.get(trace_id, {}).get("physical_hbm_bytes_written_run_stats") for trace_id in trace_ids] + ) + stat_rows.append( + { + "benchmark": benchmark, + "category": category, + "runs": len(rows), + "cycles_min": cycle_stats["min"], + "cycles_median": cycle_stats["median"], + "cycles_mean": cycle_stats["mean"], + "cycles_p95": cycle_stats["p95"], + "cycles_max": cycle_stats["max"], + "routing_tax_cycles_min": tax_stats["min"], + "routing_tax_cycles_median": tax_stats["median"], + "routing_tax_cycles_mean": tax_stats["mean"], + "routing_tax_cycles_p95": tax_stats["p95"], + "routing_tax_cycles_max": tax_stats["max"], + "hbm_read_bytes_min": read_stats["min"], + "hbm_read_bytes_median": read_stats["median"], + "hbm_read_bytes_mean": read_stats["mean"], + "hbm_read_bytes_p95": read_stats["p95"], + "hbm_read_bytes_max": read_stats["max"], + "hbm_written_bytes_min": write_stats["min"], + "hbm_written_bytes_median": write_stats["median"], + "hbm_written_bytes_mean": write_stats["mean"], + "hbm_written_bytes_p95": write_stats["p95"], + "hbm_written_bytes_max": write_stats["max"], + "measurement_note": MEASUREMENT_NOTE, + } + ) + write_csv(out_root / "gpqa_bfcl_emulation_distribution_stats.csv", stat_rows) + + +def _collect_manifest_counts(manifest: dict[str, Any]) -> dict[str, int]: + return dict(Counter(row.get("status", "unknown") for row in manifest.get("runs", []))) + + +def _routing_sample_log_stats(path: Path) -> dict[str, Any]: + if not path.exists(): + return { + "path": str(path), + "rows": 0, + "status_counts": {}, + "seconds": _stats([]), + } + rows = list(iter_jsonl(path)) + seconds = [row.get("seconds") for row in rows if row.get("status") == "passed"] + return { + "path": str(path), + "rows": len(rows), + "status_counts": dict(Counter(str(row.get("status", "unknown")) for row in rows)), + "seconds": _stats(seconds), + } + + +def _write_cost_report( + args: argparse.Namespace, + out_root: Path, + layers: list[int], + routing_summaries: list[dict[str, Any]], + replay_manifest: dict[str, Any] | None, + export_payload: dict[str, Any] | None, + prune_summary: dict[str, Any] | None, +) -> dict[str, Any]: + routing_by_benchmark = {row["benchmark"]: row for row in routing_summaries} + trace_files = sorted((out_root / "route_traces").glob("*.json")) + replay_runs = replay_manifest.get("runs", []) if replay_manifest else [] + replay_counts = _collect_manifest_counts(replay_manifest or {}) + routing_sample_stats = _routing_sample_log_stats(_routing_sample_log_path(out_root)) + replay_host_seconds = [ + float(row["host_batch_seconds"]) + for row in replay_runs + if row.get("status") == "passed" and row.get("host_batch_seconds") is not None + ] + replay_pruned_bytes = sum(int(row.get("pruned_bytes") or 0) for row in replay_runs) + replay_pruned_files = sum(int(row.get("pruned_file_count") or 0) for row in replay_runs) + observed_processed = sum(int(row.get("processed_samples", 0)) for row in routing_summaries) + completed_by_benchmark = _completed_routing_samples(out_root, args, layers) + completed_total = sum(completed_by_benchmark.values()) + full_samples = sum(FULL_SAMPLE_COUNTS.get(benchmark, 0) for benchmark in args.benchmarks) + full_runs = full_samples * len(layers) + route_wall_seconds = sum(float(row.get("orchestrator_wall_seconds", 0.0)) for row in routing_summaries) + replay_wall_seconds = float((replay_manifest or {}).get("orchestrator_wall_seconds", 0.0)) + total_observed_wall = route_wall_seconds + replay_wall_seconds + avg_route_seconds_per_processed_sample = ( + route_wall_seconds / float(observed_processed) if observed_processed else None + ) + avg_replay_seconds_per_passed_run = ( + statistics.mean(replay_host_seconds) if replay_host_seconds else None + ) + projected_route_seconds = ( + avg_route_seconds_per_processed_sample * float(full_samples) + if avg_route_seconds_per_processed_sample is not None + else None + ) + projected_replay_seconds = ( + avg_replay_seconds_per_passed_run * float(full_runs) + if avg_replay_seconds_per_passed_run is not None + else None + ) + projected_total_seconds = ( + projected_route_seconds + projected_replay_seconds + if projected_route_seconds is not None and projected_replay_seconds is not None + else None + ) + disk_sizes = { + "true_routing_bytes": _dir_size_bytes(out_root / "true_routing"), + "route_traces_bytes": _dir_size_bytes(out_root / "route_traces"), + "trace_replay_bytes": _dir_size_bytes(out_root / "trace_replay"), + "total_out_root_bytes": _dir_size_bytes(out_root), + } + observed_runs = len([row for row in replay_runs if row.get("status") in ("passed", "skipped_existing")]) + projected_total_bytes = ( + (disk_sizes["total_out_root_bytes"] / float(observed_runs)) * float(full_runs) + if observed_runs + else None + ) + report = { + "schema_version": 1, + "out_root": str(out_root), + "model_key": args.model_key, + "model_name": MODEL_CONFIGS[args.model_key]["name"], + "benchmarks": args.benchmarks, + "layers": layers, + "phase": args.write_phases, + "decode_steps": args.decode_steps, + "measurement_note": MEASUREMENT_NOTE, + "routing_summaries": routing_by_benchmark, + "completed_true_routing_samples_by_benchmark": completed_by_benchmark, + "completed_true_routing_samples": completed_total, + "route_trace_count": len(trace_files), + "route_trace_dir": str(out_root / "route_traces"), + "replay_manifest": str(out_root / "trace_replay_batch_manifest.json"), + "replay_status_counts": replay_counts, + "exported_runs": len((export_payload or {}).get("runs", [])), + "routing_failure_log": str(out_root / "failures" / "routing_failures.jsonl"), + "routing_sample_log": str(_routing_sample_log_path(out_root)), + "routing_sample_log_stats": routing_sample_stats, + "replay_failure_log": str(out_root / "failures" / "replay_failures.jsonl"), + "pruned_artifacts": prune_summary + or { + "pruned_file_count": 0, + "pruned_bytes": 0, + "note": "Artifact pruning was not requested.", + }, + "replay_pruned_artifacts": { + "pruned_file_count": replay_pruned_files, + "pruned_bytes": replay_pruned_bytes, + "note": "Artifacts pruned immediately after successful replay runs.", + }, + "observed": { + "processed_samples_this_run": observed_processed, + "completed_true_routing_samples": completed_total, + "route_wall_seconds": route_wall_seconds, + "replay_wall_seconds": replay_wall_seconds, + "total_wall_seconds": total_observed_wall, + "avg_route_seconds_per_processed_sample": avg_route_seconds_per_processed_sample, + "avg_replay_seconds_per_passed_run": avg_replay_seconds_per_passed_run, + }, + "full_scale_projection": { + "full_samples": full_samples, + "full_runs": full_runs, + "projected_route_seconds": projected_route_seconds, + "projected_replay_seconds": projected_replay_seconds, + "projected_total_seconds": projected_total_seconds, + "projected_total_hours": projected_total_seconds / 3600.0 if projected_total_seconds is not None else None, + "projected_total_output_bytes": projected_total_bytes, + }, + "disk_sizes": disk_sizes, + "known_limits": [ + "P3 reuses P2 fixed-route replay semantics.", + "Numbers are self-consistent upper bounds pending RTL calibration.", + "Replay is decode tok1 only.", + "Route replay uses zero-activation timing shape, as in P2.", + "Current P2 true-routing generator supports qwen3 only.", + ], + } + report_stem = args.report_name + write_json(out_root / f"{report_stem}.json", report) + _write_cost_report_md(out_root / f"{report_stem}.md", report) + return report + + +def _fmt_seconds(value: float | None) -> str: + if value is None: + return "n/a" + return f"{value:.1f}s ({value / 3600.0:.2f}h)" + + +def _fmt_bytes(value: float | int | None) -> str: + if value is None: + return "n/a" + value_f = float(value) + units = ["B", "KiB", "MiB", "GiB", "TiB"] + idx = 0 + while value_f >= 1024.0 and idx < len(units) - 1: + value_f /= 1024.0 + idx += 1 + return f"{value_f:.2f} {units[idx]}" + + +def _write_cost_report_md(path: Path, report: dict[str, Any]) -> None: + observed = report["observed"] + projection = report["full_scale_projection"] + disk = report["disk_sizes"] + pruned = report["pruned_artifacts"] + replay_pruned = report["replay_pruned_artifacts"] + replay_counts = report["replay_status_counts"] + routing_sample_stats = report["routing_sample_log_stats"] + lines = [ + "# Window 1 P3 Scale Batch Report", + "", + f"- model: `{report['model_name']}` (`{report['model_key']}`)", + f"- benchmarks: `{','.join(report['benchmarks'])}`", + f"- layers: `{','.join(str(x) for x in report['layers'])}`", + f"- phase: `{report['phase']}`, decode_steps: `{report['decode_steps']}`", + f"- measurement note: {report['measurement_note']}", + "", + "## Observed Batch", + "", + f"- processed true-routing samples this run: `{observed['processed_samples_this_run']}`", + f"- completed true-routing samples in output: `{observed['completed_true_routing_samples']}`", + f"- cumulative routing sample log rows: `{routing_sample_stats['rows']}`", + f"- routing sample status counts: `{routing_sample_stats['status_counts']}`", + f"- routing sample seconds median/mean/p95: " + f"`{routing_sample_stats['seconds']['median']}` / " + f"`{routing_sample_stats['seconds']['mean']}` / " + f"`{routing_sample_stats['seconds']['p95']}`", + f"- route trace files: `{report['route_trace_count']}`", + f"- replay status counts: `{replay_counts}`", + f"- exported timing runs: `{report['exported_runs']}`", + f"- replay-time pruned artifacts: `{replay_pruned['pruned_file_count']}` files, " + f"`{_fmt_bytes(replay_pruned['pruned_bytes'])}`", + f"- pruned heavy artifacts: `{pruned['pruned_file_count']}` files, `{_fmt_bytes(pruned['pruned_bytes'])}`", + f"- route generation wall time: `{_fmt_seconds(observed['route_wall_seconds'])}`", + f"- replay wall time: `{_fmt_seconds(observed['replay_wall_seconds'])}`", + f"- total observed wall time: `{_fmt_seconds(observed['total_wall_seconds'])}`", + f"- avg route seconds/sample: `{observed['avg_route_seconds_per_processed_sample']}`", + f"- avg replay seconds/passed run: `{observed['avg_replay_seconds_per_passed_run']}`", + "", + "## Full-Scale Projection", + "", + f"- full samples: `{projection['full_samples']}`", + f"- full representative-layer runs: `{projection['full_runs']}`", + f"- projected route generation: `{_fmt_seconds(projection['projected_route_seconds'])}`", + f"- projected replay: `{_fmt_seconds(projection['projected_replay_seconds'])}`", + f"- projected total: `{_fmt_seconds(projection['projected_total_seconds'])}`", + f"- projected output size: `{_fmt_bytes(projection['projected_total_output_bytes'])}`", + "", + "## Disk", + "", + f"- true routing: `{_fmt_bytes(disk['true_routing_bytes'])}`", + f"- route traces: `{_fmt_bytes(disk['route_traces_bytes'])}`", + f"- trace replay: `{_fmt_bytes(disk['trace_replay_bytes'])}`", + f"- total output root: `{_fmt_bytes(disk['total_out_root_bytes'])}`", + "", + "## Limits", + "", + ] + lines.extend(f"- {item}" for item in report["known_limits"]) + lines.extend( + [ + "", + "## Failure Logs", + "", + f"- routing: `{report['routing_failure_log']}`", + f"- replay: `{report['replay_failure_log']}`", + "", + ] + ) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines), encoding="utf-8") + + +def run(args: argparse.Namespace) -> dict[str, Any]: + ensure_paths() + out_root = args.out_root + out_root.mkdir(parents=True, exist_ok=True) + (out_root / "failures").mkdir(parents=True, exist_ok=True) + (out_root / "failures" / "routing_failures.jsonl").touch(exist_ok=True) + (out_root / "failures" / "replay_failures.jsonl").touch(exist_ok=True) + layers = _parse_csv_ints(args.layers) + routing_summaries: list[dict[str, Any]] = [] + + if not args.skip_routing: + for benchmark in args.benchmarks: + routing_summaries.append(_run_true_routing(args, benchmark, layers, out_root)) + else: + for benchmark in args.benchmarks: + summary_path = _routing_summary_path(out_root, benchmark, args.model_key) + if summary_path.exists(): + routing_summaries.append(load_json(summary_path)) + + if not args.skip_build_traces: + for benchmark in args.benchmarks: + _build_route_traces(args, benchmark, layers, out_root) + + replay_manifest = None + if not args.skip_replay: + replay_manifest = _run_replay(args, out_root) + else: + manifest_path = out_root / "trace_replay_batch_manifest.json" + if manifest_path.exists(): + replay_manifest = load_json(manifest_path) + + export_payload = None + if not args.skip_export: + export_payload = _export(args, out_root) + else: + export_path = out_root / "gpqa_bfcl_emulation_timing_summary.json" + if export_path.exists(): + export_payload = load_json(export_path) + + prune_summary = _prune_heavy_artifacts(out_root) if args.prune_heavy_artifacts else None + + report = _write_cost_report( + args, + out_root, + layers, + routing_summaries, + replay_manifest, + export_payload, + prune_summary, + ) + write_json(out_root / "p3_scale_batch_summary.json", report) + _append_jsonl( + out_root / "p3_scale_batch_history.jsonl", + { + "timestamp": time.time(), + "benchmarks": args.benchmarks, + "bfcl_limit": args.bfcl_limit, + "gpqa_limit": args.gpqa_limit, + "layers": layers, + "write_phases": args.write_phases, + "decode_steps": args.decode_steps, + "skip_routing": args.skip_routing, + "skip_build_traces": args.skip_build_traces, + "skip_replay": args.skip_replay, + "skip_export": args.skip_export, + "prune_heavy_artifacts": args.prune_heavy_artifacts, + "observed": report["observed"], + "replay_status_counts": report["replay_status_counts"], + "exported_runs": report["exported_runs"], + "replay_pruned_artifacts": { + "pruned_file_count": report["replay_pruned_artifacts"].get("pruned_file_count"), + "pruned_bytes": report["replay_pruned_artifacts"].get("pruned_bytes"), + }, + "pruned_artifacts": { + "pruned_file_count": report["pruned_artifacts"].get("pruned_file_count"), + "pruned_bytes": report["pruned_artifacts"].get("pruned_bytes"), + }, + }, + ) + return report + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT) + parser.add_argument("--model-key", choices=sorted(MODEL_CONFIGS), default="qwen3") + parser.add_argument("--benchmarks", type=_parse_csv_strings, default=_parse_csv_strings("bfcl_v3,gpqa_diamond")) + parser.add_argument("--bfcl-limit", type=int, default=50) + parser.add_argument("--gpqa-limit", type=int, default=20) + parser.add_argument("--layers", default=DEFAULT_LAYERS) + parser.add_argument("--write-phases", choices=("prefill", "decode", "both"), default="decode") + parser.add_argument("--decode-steps", type=int, default=1) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--max-input-tokens", type=int) + parser.add_argument("--sort-by-length", action="store_true") + parser.add_argument("--mlen", type=int, default=128) + parser.add_argument("--blen", type=int, default=4) + parser.add_argument("--emu-threads", type=int, default=1) + parser.add_argument("--cwd", type=Path, default=REPO_ROOT) + parser.add_argument("--replay-limit", type=int) + parser.add_argument("--stage-profile", action="store_true") + parser.add_argument("--resume", action="store_true") + parser.add_argument("--skip-existing", action="store_true") + parser.add_argument("--keep-going", action="store_true") + parser.add_argument("--keep-dumps", action="store_true") + parser.add_argument("--skip-routing", action="store_true") + parser.add_argument("--skip-build-traces", action="store_true") + parser.add_argument("--skip-replay", action="store_true") + parser.add_argument("--skip-export", action="store_true") + parser.add_argument("--prune-heavy-artifacts", action="store_true") + parser.add_argument("--report-name", default="s1_cost_report") + args = parser.parse_args() + + unknown = [benchmark for benchmark in args.benchmarks if benchmark not in INPUT_FILES] + if unknown: + raise ValueError(f"unknown benchmarks: {unknown}") + report = run(args) + print(f"wrote {args.out_root / (args.report_name + '.md')}") + print( + "processed_samples=" + f"{report['observed']['processed_samples_this_run']} " + f"completed_samples={report['observed']['completed_true_routing_samples']} " + f"route_traces={report['route_trace_count']} " + f"exported_runs={report['exported_runs']}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p3/wait_p3_rev_stage_b.py b/transactional_emulator/testbench/window1_p3/wait_p3_rev_stage_b.py new file mode 100644 index 00000000..4a9d46f0 --- /dev/null +++ b/transactional_emulator/testbench/window1_p3/wait_p3_rev_stage_b.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 +"""Wait for P3-rev Stage A gates, then launch Stage B replay/export. + +This is intended to run inside a detached tmux session. It does not change +timing semantics; it only polls checkpoint files and starts the existing +P3-rev sidecar once route traces and determinism evidence are ready. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import signal +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import PLENA_ROOT + + +DEFAULT_OUT_ROOT = PLENA_ROOT / "outputs" / "window1_p3" +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def _load_json(path: Path) -> dict[str, Any] | None: + if not path.exists(): + return None + with path.open(encoding="utf-8") as handle: + return json.load(handle) + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f"{path.name}.tmp.{time.time_ns()}") + tmp.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + tmp.replace(path) + + +def _line_count(path: Path) -> int: + if not path.exists(): + return 0 + with path.open(encoding="utf-8") as handle: + return sum(1 for _ in handle) + + +def _tmux_session_exists(name: str) -> bool: + if shutil.which("tmux") is None: + return False + return subprocess.run( + ["tmux", "has-session", "-t", name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode == 0 + + +def _kill_tmux_session(name: str) -> int: + if shutil.which("tmux") is None: + return 127 + pane_pids = _tmux_pane_pids(name) + descendant_pids = _descendant_pids(pane_pids) + rc = subprocess.run( + ["tmux", "kill-session", "-t", name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode + _terminate_pids(descendant_pids) + return rc + + +def _tmux_pane_pids(name: str) -> list[int]: + if shutil.which("tmux") is None: + return [] + result = subprocess.run( + ["tmux", "list-panes", "-t", name, "-F", "#{pane_pid}"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + if result.returncode != 0: + return [] + pids = [] + for line in result.stdout.splitlines(): + try: + pids.append(int(line.strip())) + except ValueError: + pass + return pids + + +def _descendant_pids(root_pids: list[int]) -> list[int]: + if not root_pids: + return [] + result = subprocess.run( + ["ps", "-eo", "pid=,ppid="], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + if result.returncode != 0: + return [] + children: dict[int, list[int]] = {} + for line in result.stdout.splitlines(): + parts = line.split() + if len(parts) != 2: + continue + try: + pid = int(parts[0]) + ppid = int(parts[1]) + except ValueError: + continue + children.setdefault(ppid, []).append(pid) + roots = set(root_pids) + seen: set[int] = set() + stack = list(root_pids) + while stack: + pid = stack.pop() + for child in children.get(pid, []): + if child in seen or child in roots: + continue + seen.add(child) + stack.append(child) + return sorted(seen, reverse=True) + + +def _terminate_pids(pids: list[int]) -> None: + for sig in (signal.SIGTERM, signal.SIGKILL): + live = [] + for pid in pids: + try: + os.kill(pid, 0) + except OSError: + continue + live.append(pid) + if not live: + return + for pid in live: + try: + os.kill(pid, sig) + except OSError: + pass + time.sleep(1) + + +def _run_progress(args: argparse.Namespace) -> None: + cmd = [ + sys.executable, + "-m", + "transactional_emulator.testbench.window1_p3.run_p3_rev_subset", + "--out-root", + str(args.out_root), + "--actions", + "build,progress,report", + "--stage-profile", + "--skip-existing", + "--keep-going", + "--prune-success-artifacts", + "--no-stage-status", + ] + subprocess.run(cmd, cwd=args.cwd, check=True) + + +def _ready(args: argparse.Namespace) -> tuple[bool, dict[str, Any]]: + _run_progress(args) + progress = _load_json(args.out_root / "p3_rev_progress_snapshot.json") or {} + determinism = _load_json(args.out_root / "p3_rev_determinism.json") + failure_counts = { + path.name: _line_count(path) + for path in sorted((args.out_root / "failures").glob("*.jsonl")) + } + expected = int(progress.get("expected_representative_runs") or 0) + traces = int(progress.get("selected_trace_files") or 0) + route_ready = expected > 0 and traces >= expected + determinism_ready = ( + (not args.require_determinism) + or (determinism is not None and determinism.get("status") == "passed") + ) + failure_free = all(count == 0 for count in failure_counts.values()) + + stopped_sessions: dict[str, int] = {} + if route_ready: + for name in args.stop_session_on_route_ready: + if _tmux_session_exists(name): + stopped_sessions[name] = _kill_tmux_session(name) + if stopped_sessions: + time.sleep(2) + + blocking_sessions = [name for name in args.wait_session if _tmux_session_exists(name)] + sessions_ready = not blocking_sessions + state = { + "schema_version": 1, + "updated_at_unix": time.time(), + "expected_representative_runs": expected, + "selected_trace_files": traces, + "route_ready": route_ready, + "determinism_status": None if determinism is None else determinism.get("status"), + "determinism_ready": determinism_ready, + "failure_counts": failure_counts, + "failure_free": failure_free, + "stopped_sessions_on_route_ready": stopped_sessions, + "blocking_sessions": blocking_sessions, + "sessions_ready": sessions_ready, + "ready": route_ready and determinism_ready and failure_free and sessions_ready, + } + return bool(state["ready"]), state + + +def _launch_stage_b(args: argparse.Namespace) -> int: + cmd = [ + sys.executable, + "-m", + "transactional_emulator.testbench.window1_p3.run_p3_rev_subset", + "--out-root", + str(args.out_root), + "--actions", + "replay,export,report", + "--stage-profile", + "--skip-existing", + "--keep-going", + "--prune-success-artifacts", + "--width", + str(args.width), + "--report-name", + "p3_rev_report", + ] + print("launching Stage B:", " ".join(cmd), flush=True) + return subprocess.run(cmd, cwd=args.cwd).returncode + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out-root", type=Path, default=DEFAULT_OUT_ROOT) + parser.add_argument("--cwd", type=Path, default=REPO_ROOT) + parser.add_argument("--poll-seconds", type=int, default=600) + parser.add_argument("--width", type=int, default=8) + parser.add_argument("--require-determinism", action="store_true") + parser.add_argument("--wait-session", action="append", default=["plena_p3_rev_route_build", "plena_p3_rev_gpqa_route"]) + parser.add_argument( + "--stop-session-on-route-ready", + action="append", + default=[], + help=( + "tmux session to stop once all route traces are present. " + "Use this for low-priority partial replay sessions so final Stage B " + "can safely take over with skip-existing instead of waiting for the " + "partial queue to drain." + ), + ) + parser.add_argument("--max-polls", type=int) + args = parser.parse_args() + + status_path = args.out_root / "p3_rev_stage_b_waiter_status.json" + poll = 0 + while True: + poll += 1 + try: + ready, state = _ready(args) + except Exception as exc: + state = { + "schema_version": 1, + "updated_at_unix": time.time(), + "ready": False, + "error_type": type(exc).__name__, + "error": str(exc), + } + ready = False + state["poll"] = poll + state["poll_seconds"] = args.poll_seconds + _write_json(status_path, state) + print(json.dumps(state, sort_keys=True), flush=True) + if ready: + state["stage_b_started_at_unix"] = time.time() + _write_json(status_path, state) + rc = _launch_stage_b(args) + state["stage_b_return_code"] = rc + state["stage_b_ended_at_unix"] = time.time() + _write_json(status_path, state) + return rc + if args.max_polls is not None and poll >= args.max_polls: + return 2 + time.sleep(args.poll_seconds) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p3/watch_p3_rev_raw_build.py b/transactional_emulator/testbench/window1_p3/watch_p3_rev_raw_build.py new file mode 100644 index 00000000..bb873de3 --- /dev/null +++ b/transactional_emulator/testbench/window1_p3/watch_p3_rev_raw_build.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Lightweight P3-rev bridge from completed raw routing rows to route traces. + +This watcher does not generate routing and does not run replay. It only notices +when true_routing JSONL rows get ahead of built route_trace JSON files, then +delegates to run_p3_rev_subset's existing build/progress/report path. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import time +from pathlib import Path + + +ROUTE_TRACE_RE = re.compile(r"^qwen3_(bfcl_v3|gpqa_diamond)_(.*)_l(0|12|23)_decode_d1\.json$") + + +def _selected(out_root: Path) -> dict[str, set[str]]: + manifest = json.loads((out_root / "p3_rev_selected_samples.json").read_text()) + return {bench: {str(row["sample_id"]) for row in obj["selected_samples"]} for bench, obj in manifest["benchmarks"].items()} + + +def _counts(out_root: Path, selected: dict[str, set[str]]) -> tuple[int, int, int, int]: + built_seen: dict[str, dict[str, set[int]]] = {bench: {} for bench in selected} + for path in (out_root / "route_traces").glob("*.json"): + match = ROUTE_TRACE_RE.match(path.name) + if not match: + continue + bench, sample_id, layer = match.group(1), match.group(2), int(match.group(3)) + if bench in selected and sample_id in selected[bench]: + built_seen[bench].setdefault(sample_id, set()).add(layer) + + raw = set() + for path in (out_root / "true_routing").glob("*.jsonl"): + for line in path.read_text(errors="ignore").splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue + bench = str(row.get("benchmark", "")) + sample_id = str(row.get("sample_id", "")) + if ( + bench in selected + and sample_id in selected[bench] + and row.get("phase") == "decode" + and int(row.get("decode_step", 1)) == 1 + ): + raw.add((bench, sample_id, int(row["layer"]))) + + built = { + (bench, sample_id, layer) + for bench, sample_layers in built_seen.items() + for sample_id, layers in sample_layers.items() + for layer in layers + } + expected = sum(len(samples) * 3 for samples in selected.values()) + return len(raw), len(built), len(raw - built), expected + + +def _run_incremental_build(repo: Path, out_root: Path, report_name: str) -> int: + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = "" + env["PYTHONPATH"] = ":".join( + [ + str(repo), + str(repo / "PLENA_Compiler"), + str(repo / "PLENA_Tools"), + str(repo / "transactional_emulator" / "testbench"), + env.get("PYTHONPATH", ""), + ] + ) + cmd = [ + os.environ.get("PYTHON", "python"), + "-m", + "transactional_emulator.testbench.window1_p3.run_p3_rev_subset", + "--out-root", + str(out_root), + "--actions", + "build,progress,report", + "--stage-profile", + "--skip-existing", + "--keep-going", + "--prune-success-artifacts", + "--no-stage-status", + "--report-name", + report_name, + ] + return subprocess.run(cmd, cwd=repo, env=env, check=False).returncode + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", type=Path, default=Path.cwd()) + parser.add_argument("--out-root", type=Path, required=True) + parser.add_argument("--poll-seconds", type=int, default=120) + parser.add_argument("--report-name", default="p3_rev_incremental_build_report") + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + + selected = _selected(args.out_root) + poll = 0 + while True: + poll += 1 + raw_count, built_count, raw_minus_built, expected = _counts(args.out_root, selected) + print( + json.dumps( + { + "schema_version": 1, + "poll": poll, + "raw_selected_keys": raw_count, + "built_selected_keys": built_count, + "raw_minus_built": raw_minus_built, + "expected": expected, + "updated_at_unix": time.time(), + }, + sort_keys=True, + ), + flush=True, + ) + if raw_minus_built: + rc = _run_incremental_build(args.repo, args.out_root, args.report_name) + print(json.dumps({"incremental_build_returncode": rc, "updated_at_unix": time.time()}), flush=True) + if args.once or built_count >= expected: + return 0 + time.sleep(args.poll_seconds) + + +if __name__ == "__main__": + raise SystemExit(main())