Files
cs249r_book/book/tools/scripts/audit/math_chapter_workflow.py
2026-05-20 18:56:52 -07:00

682 lines
25 KiB
Python

#!/usr/bin/env python3
"""Build MLSysBook chapter packets for math-accuracy audit rounds.
The script uses the existing `book/binder` PDF build path. For each selected
unit it builds a standalone PDF, copies the generated PDF and keep-tex output
into a stable packet directory, and writes small-batch prompts for Codex agents.
"""
from __future__ import annotations
import argparse
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable
try:
import yaml
except ImportError as exc: # pragma: no cover - dependency is declared by repo
raise SystemExit("PyYAML is required. Install project requirements first.") from exc
REPO_ROOT = Path(__file__).resolve().parents[4]
BOOK_DIR = REPO_ROOT / "book" / "quarto"
BINDER = REPO_ROOT / "book" / "binder"
RUNS_DIR = REPO_ROOT / "review" / "math-audit" / "runs"
PDF_NAMES = {
"vol1": "Machine-Learning-Systems-Vol1.pdf",
"vol2": "Machine-Learning-Systems-Vol2.pdf",
}
TEX_NAMES = {
"vol1": "Machine-Learning-Systems-Vol1.tex",
"vol2": "Machine-Learning-Systems-Vol2.tex",
}
SKIP_STEMS = {"index", "references"}
@dataclass(frozen=True)
class BuildUnit:
volume: str
config_position: int
stem: str
rel_path: str
source_path: Path
kind: str
source_section: str
@property
def packet_name(self) -> str:
return f"{self.config_position:02d}-{self.stem}"
@property
def chapter_id(self) -> str:
return f"{self.volume}/{self.stem}"
def now_iso() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
def slugify(value: str) -> str:
return re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-")
def parse_csv(value: str | None, default: list[str]) -> list[str]:
if not value:
return default
return [part.strip() for part in value.split(",") if part.strip()]
def read_yaml(path: Path) -> dict[str, Any]:
return yaml.safe_load(path.read_text(encoding="utf-8")) or {}
def write_yaml(path: Path, data: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=False), encoding="utf-8")
def rel_entry_path(entry: Any) -> str | None:
if isinstance(entry, str):
return entry
if isinstance(entry, dict):
file_value = entry.get("file")
return str(file_value) if file_value else None
return None
def classify_unit(rel_path: str, source_section: str) -> str:
stem = Path(rel_path).stem
parts = Path(rel_path).parts
if source_section == "appendices":
return "glossary" if stem == "glossary" else "appendix"
if "frontmatter" in parts:
return "frontmatter"
if "parts" in parts:
return "part"
if "backmatter" in parts:
if stem.startswith("appendix_"):
return "appendix"
if stem == "glossary":
return "glossary"
return "backmatter"
return "chapter"
def discover_units(volumes: Iterable[str]) -> list[BuildUnit]:
units: list[BuildUnit] = []
for volume in volumes:
config_path = BOOK_DIR / "config" / f"_quarto-pdf-{volume}.yml"
if not config_path.exists():
raise FileNotFoundError(f"Missing PDF config: {config_path}")
config = read_yaml(config_path)
book_config = config.get("book", {})
position = 0
for section_name in ("chapters", "appendices"):
for entry in book_config.get(section_name, []) or []:
rel_path = rel_entry_path(entry)
if not rel_path:
continue
stem = Path(rel_path).stem
if stem in SKIP_STEMS:
continue
source_path = BOOK_DIR / rel_path
if not source_path.exists():
raise FileNotFoundError(f"Configured source does not exist: {source_path}")
position += 1
units.append(
BuildUnit(
volume=volume,
config_position=position,
stem=stem,
rel_path=rel_path,
source_path=source_path,
kind=classify_unit(rel_path, section_name),
source_section=section_name,
)
)
return units
def include_unit(unit: BuildUnit, args: argparse.Namespace) -> bool:
if args.include_all_units:
return True
if unit.kind == "frontmatter" and not args.include_frontmatter:
return False
if unit.kind == "part" and not args.include_parts:
return False
if unit.kind == "glossary" and not args.include_glossary:
return False
if unit.kind == "backmatter" and not args.include_backmatter:
return False
return True
def resolve_only_specs(specs: list[str], units: list[BuildUnit]) -> set[tuple[str, str]]:
resolved: set[tuple[str, str]] = set()
for raw_spec in specs:
spec = raw_spec.strip()
if not spec:
continue
if ":" in spec:
volume, stem = spec.split(":", 1)
elif "/" in spec:
volume, stem = spec.split("/", 1)
else:
matches = [unit for unit in units if unit.stem == spec]
if not matches:
raise ValueError(f"No unit named {spec!r}")
if len(matches) > 1:
choices = ", ".join(unit.chapter_id for unit in matches)
raise ValueError(f"Ambiguous unit {spec!r}; use one of: {choices}")
volume, stem = matches[0].volume, matches[0].stem
resolved.add((volume.strip(), stem.strip().removesuffix(".qmd")))
return resolved
def selected_units(args: argparse.Namespace) -> list[BuildUnit]:
volumes = parse_csv(args.volumes, ["vol1", "vol2"])
units = [unit for unit in discover_units(volumes) if include_unit(unit, args)]
if args.only:
specs = resolve_only_specs(parse_csv(args.only, []), units)
units = [unit for unit in units if (unit.volume, unit.stem) in specs]
if args.limit is not None:
units = units[: args.limit]
return units
def output_dir_for(volume: str) -> Path:
return BOOK_DIR / "_build" / f"pdf-{volume}"
def clean_previous_build(volume: str) -> None:
output_dir = output_dir_for(volume)
if output_dir.exists():
shutil.rmtree(output_dir)
tex_file = BOOK_DIR / TEX_NAMES[volume]
if tex_file.exists():
tex_file.unlink()
def run_command(cmd: list[str], cwd: Path, log_path: Path, verbose: bool) -> int:
log_path.parent.mkdir(parents=True, exist_ok=True)
env = os.environ.copy()
process = subprocess.Popen(
cmd,
cwd=str(cwd),
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
assert process.stdout is not None
with log_path.open("w", encoding="utf-8") as log:
for line in process.stdout:
log.write(line)
if verbose:
print(line, end="")
return process.wait()
def build_unit(unit: BuildUnit, run_dir: Path, force: bool, verbose: bool) -> dict[str, Any]:
packet_dir = run_dir / "packets" / unit.volume / unit.packet_name
if packet_dir.exists() and force:
shutil.rmtree(packet_dir)
packet_dir.mkdir(parents=True, exist_ok=True)
log_path = packet_dir / "build.log"
clean_previous_build(unit.volume)
cmd = [str(BINDER), "build", "pdf", unit.stem, f"--{unit.volume}"]
if verbose:
cmd.append("-v")
print(f"Building {unit.chapter_id} -> {packet_dir}", flush=True)
return_code = run_command(cmd, REPO_ROOT, log_path, verbose=verbose)
if return_code != 0:
return {
"chapter": unit.chapter_id,
"kind": unit.kind,
"packet_dir": str(packet_dir),
"build_log": str(log_path),
"status": "build_failed",
"return_code": return_code,
}
return package_unit(unit, run_dir, packet_dir, log_path)
def build_volume(volume: str, run_dir: Path, force: bool, verbose: bool) -> dict[str, Any]:
dest_dir = run_dir / "volumes" / volume
if dest_dir.exists() and force:
shutil.rmtree(dest_dir)
dest_dir.mkdir(parents=True, exist_ok=True)
log_path = dest_dir / "build.log"
clean_previous_build(volume)
cmd = [str(BINDER), "build", "pdf", f"--{volume}"]
if verbose:
cmd.append("-v")
print(f"Building full {volume} -> {dest_dir}", flush=True)
return_code = run_command(cmd, REPO_ROOT, log_path, verbose=verbose)
if return_code != 0:
return {
"volume": volume,
"dest_dir": str(dest_dir),
"build_log": str(log_path),
"status": "build_failed",
"return_code": return_code,
}
pdf_src = output_dir_for(volume) / PDF_NAMES[volume]
tex_src = BOOK_DIR / TEX_NAMES[volume]
if not pdf_src.exists() or not tex_src.exists():
missing = [str(path) for path in (pdf_src, tex_src) if not path.exists()]
return {
"volume": volume,
"dest_dir": str(dest_dir),
"build_log": str(log_path),
"status": "missing_artifact",
"missing": missing,
}
pdf_dest = dest_dir / f"{volume}-full.pdf"
tex_dest = dest_dir / f"{volume}-full.tex"
shutil.copy2(pdf_src, pdf_dest)
shutil.copy2(tex_src, tex_dest)
math_context = dest_dir / "math_context.md"
extract_math_context(tex_dest, math_context)
record = {
"volume": volume,
"dest_dir": str(dest_dir),
"status": "built",
"pdf": str(pdf_dest),
"tex": str(tex_dest),
"math_context": str(math_context),
"build_log": str(log_path),
"built_at": now_iso(),
}
write_yaml(dest_dir / "manifest.yml", record)
return record
def package_unit(unit: BuildUnit, run_dir: Path, packet_dir: Path, log_path: Path) -> dict[str, Any]:
pdf_src = output_dir_for(unit.volume) / PDF_NAMES[unit.volume]
tex_src = BOOK_DIR / TEX_NAMES[unit.volume]
missing = [str(path) for path in (pdf_src, tex_src, unit.source_path) if not path.exists()]
if missing:
return {
"chapter": unit.chapter_id,
"kind": unit.kind,
"packet_dir": str(packet_dir),
"build_log": str(log_path),
"status": "missing_artifact",
"missing": missing,
}
artifact_stem = f"{unit.volume}-{unit.stem}"
pdf_dest = packet_dir / f"{artifact_stem}.pdf"
tex_dest = packet_dir / f"{artifact_stem}.tex"
qmd_dest = packet_dir / f"{artifact_stem}.qmd"
shutil.copy2(pdf_src, pdf_dest)
shutil.copy2(tex_src, tex_dest)
shutil.copy2(unit.source_path, qmd_dest)
for optional_name in ("FIGURE_LIST.txt", "index_figures.txt"):
optional_src = output_dir_for(unit.volume) / optional_name
if optional_src.exists():
shutil.copy2(optional_src, packet_dir / optional_name)
math_context = packet_dir / "math_context.md"
extract_math_context(tex_dest, math_context)
audit_output = run_dir / "audits" / f"{artifact_stem}.yml"
packet_prompt = packet_dir / "agent_prompt.md"
record = {
"chapter": unit.chapter_id,
"volume": unit.volume,
"stem": unit.stem,
"kind": unit.kind,
"source_section": unit.source_section,
"source_qmd": str(unit.source_path),
"source_qmd_copy": str(qmd_dest),
"rel_path": unit.rel_path,
"packet_dir": str(packet_dir),
"status": "built",
"pdf": str(pdf_dest),
"tex": str(tex_dest),
"math_context": str(math_context),
"audit_output": str(audit_output),
"agent_prompt": str(packet_prompt),
"build_log": str(log_path),
"built_at": now_iso(),
}
write_yaml(packet_dir / "manifest.yml", record)
packet_prompt.write_text(agent_prompt_for_records([record], 1), encoding="utf-8")
return record
def extract_math_context(tex_path: Path, output_path: Path) -> None:
lines = tex_path.read_text(encoding="utf-8", errors="replace").splitlines()
env_start = re.compile(
r"\\begin\{(equation\*?|align\*?|aligned|gather\*?|multline\*?|split|flalign\*?|eqnarray\*?)\}"
)
env_end = re.compile(
r"\\end\{(equation\*?|align\*?|aligned|gather\*?|multline\*?|split|flalign\*?|eqnarray\*?)\}"
)
display_start_re = re.compile(r"(?<!\\)\\\[")
display_end_re = re.compile(r"(?<!\\)\\\]")
blocks: list[tuple[str, int, int, list[str]]] = []
inline_lines: list[tuple[int, str]] = []
i = 0
while i < len(lines):
line = lines[i]
start_match = env_start.search(line)
display_start = display_start_re.search(line) is not None
if start_match or display_start:
kind = start_match.group(1) if start_match else "display"
start = i
end = i
while end < len(lines):
if (start_match and env_end.search(lines[end])) or (
display_start and display_end_re.search(lines[end])
):
break
end += 1
end = min(end, len(lines) - 1)
context_start = max(0, start - 2)
context_end = min(len(lines), end + 3)
blocks.append((kind, start + 1, end + 1, lines[context_start:context_end]))
i = end + 1
continue
if "\\(" in line or "$" in line:
stripped = line.strip()
if stripped and not stripped.startswith("%"):
inline_lines.append((i + 1, stripped[:500]))
i += 1
out: list[str] = [
"# Math Candidate Context",
"",
"This helper is not exhaustive. Agents must still read the full TeX and QMD.",
"",
f"- source_tex: {tex_path}",
f"- display_blocks: {len(blocks)}",
f"- inline_candidate_lines: {len(inline_lines)}",
"",
]
for index, (kind, start_line, end_line, context) in enumerate(blocks, start=1):
out.extend(
[
f"## Display Block {index}",
"",
f"- kind: {kind}",
f"- tex_lines: {start_line}-{end_line}",
"",
"```tex",
*context,
"```",
"",
]
)
if inline_lines:
out.extend(["## Inline Math Candidate Lines", ""])
for line_no, text in inline_lines:
out.append(f"- line {line_no}: `{text}`")
out.append("")
output_path.write_text("\n".join(out), encoding="utf-8")
def agent_prompt_for_records(records: list[dict[str, Any]], batch_number: int) -> str:
lines = [
f"# MLSysBook Math Audit Batch {batch_number:03d}",
"",
"Use extra-high reasoning. The goal is mathematical accuracy of calculations.",
"Read each TeX file, the copied source QMD, and the PDF when rendered context matters.",
"The math_context.md file is only a locator; do not rely on it as complete coverage.",
"",
"For every numeric, algebraic, dimensional, probability, scaling, complexity,",
"latency, throughput, memory, energy, carbon, or cost calculation:",
"- recompute it independently from the stated assumptions;",
"- verify unit conversions and orders of magnitude;",
"- verify that surrounding prose matches the formula and result;",
"- flag incorrect arithmetic, inconsistent assumptions, or prose/result mismatch;",
"- do not rewrite source files in this audit pass.",
"",
"Write one YAML audit file per packet at the exact audit_output path below.",
"Use status: clear with issues: [] when no math issues are found.",
"At the top level, set chapter/source_tex/source_qmd/source_pdf to the",
"exact values for that packet.",
"",
"YAML schema:",
"```yaml",
"chapter: vol1/introduction",
"status: clear | issues_found",
"reviewed_at: '2026-05-20T00:00:00Z'",
"auditor: codex-agent",
"source_tex: /abs/path/to/file.tex",
"source_qmd: /abs/path/to/file.qmd",
"source_pdf: /abs/path/to/file.pdf",
"issues:",
" - id: vol1-introduction-MATH-001",
" severity: high | medium | low",
" tex_line: 123",
" qmd_line: 456",
" rendered_context: short description of where it appears",
" expression: formula or numeric claim being checked",
" prose_claim: prose tied to the calculation",
" observed_result: value or relationship in the book",
" expected_result: recomputed correct value or relationship",
" reasoning: concise derivation showing the correction",
" proposed_correction: exact correction to consider",
" confidence: high | medium | low",
"```",
"",
"Packets:",
]
for idx, record in enumerate(records, start=1):
lines.extend(
[
f"{idx}. {record['chapter']}",
f" packet_dir: {record['packet_dir']}",
f" tex: {record['tex']}",
f" source_qmd: {record['source_qmd_copy']}",
f" pdf: {record['pdf']}",
f" math_context: {record['math_context']}",
f" audit_output: {record['audit_output']}",
]
)
lines.append("")
return "\n".join(lines)
def write_agent_batches(run_dir: Path, packet_records: list[dict[str, Any]], batch_size: int) -> list[str]:
batch_dir = run_dir / "agent_batches"
batch_dir.mkdir(parents=True, exist_ok=True)
batch_paths: list[str] = []
built_packets = [record for record in packet_records if record.get("status") == "built"]
for start in range(0, len(built_packets), batch_size):
batch = built_packets[start : start + batch_size]
batch_number = start // batch_size + 1
path = batch_dir / f"batch-{batch_number:03d}.md"
path.write_text(agent_prompt_for_records(batch, batch_number), encoding="utf-8")
batch_paths.append(str(path))
return batch_paths
def cmd_list(args: argparse.Namespace) -> int:
for unit in selected_units(args):
print(f"{unit.chapter_id:36} {unit.kind:12} {unit.rel_path}")
return 0
def cmd_build(args: argparse.Namespace) -> int:
run_id = slugify(args.run_id or datetime.now().strftime("round-%Y%m%d-%H%M%S"))
run_dir = RUNS_DIR / run_id
run_dir.mkdir(parents=True, exist_ok=True)
(run_dir / "audits").mkdir(parents=True, exist_ok=True)
units = selected_units(args)
packet_records: list[dict[str, Any]] = []
volume_records: list[dict[str, Any]] = []
failures: list[dict[str, Any]] = []
if args.mode in ("volumes", "both"):
for volume in parse_csv(args.volumes, ["vol1", "vol2"]):
record = build_volume(volume, run_dir, args.force, args.verbose)
volume_records.append(record)
if record.get("status") != "built":
failures.append(record)
if not args.continue_on_error:
break
if args.mode in ("chapters", "both") and not (failures and not args.continue_on_error):
for unit in units:
record = build_unit(unit, run_dir, args.force, args.verbose)
packet_records.append(record)
if record.get("status") != "built":
failures.append(record)
if not args.continue_on_error:
break
batch_paths = write_agent_batches(run_dir, packet_records, args.batch_size)
run_manifest = {
"run_id": run_id,
"run_dir": str(run_dir),
"created_at": now_iso(),
"mode": args.mode,
"volumes": parse_csv(args.volumes, ["vol1", "vol2"]),
"unit_count": len(units),
"built_packet_count": len([r for r in packet_records if r.get("status") == "built"]),
"failure_count": len(failures),
"batch_size": args.batch_size,
"agent_batches": batch_paths,
"volumes_built": volume_records,
"packets": packet_records,
"failures": failures,
}
write_yaml(run_dir / "run_manifest.yml", run_manifest)
print(f"Run directory: {run_dir}")
print(f"Built packets: {run_manifest['built_packet_count']} / {len(packet_records)}")
print(f"Agent batches: {len(batch_paths)}")
if failures:
print(f"Failures: {len(failures)}")
return 1
return 0
def cmd_prompts(args: argparse.Namespace) -> int:
run_dir = Path(args.run_dir).resolve()
manifest = read_yaml(run_dir / "run_manifest.yml")
packet_records = manifest.get("packets", [])
batch_paths = write_agent_batches(run_dir, packet_records, args.batch_size)
manifest["batch_size"] = args.batch_size
manifest["agent_batches"] = batch_paths
write_yaml(run_dir / "run_manifest.yml", manifest)
print(f"Wrote {len(batch_paths)} agent batch prompts under {run_dir / 'agent_batches'}")
return 0
def cmd_summarize(args: argparse.Namespace) -> int:
run_dir = Path(args.run_dir).resolve()
audit_dir = run_dir / "audits"
manifest_path = run_dir / "run_manifest.yml"
audit_to_chapter: dict[str, str] = {}
if manifest_path.exists():
manifest = read_yaml(manifest_path)
for packet in manifest.get("packets", []) or []:
audit_output = packet.get("audit_output")
chapter = packet.get("chapter")
if audit_output and chapter:
audit_to_chapter[str(Path(audit_output).resolve())] = chapter
audit_files = sorted(audit_dir.glob("*.yml"))
audits: list[dict[str, Any]] = []
issues: list[dict[str, Any]] = []
for audit_file in audit_files:
data = read_yaml(audit_file)
inferred_chapter = audit_to_chapter.get(str(audit_file.resolve()))
if inferred_chapter and not data.get("chapter"):
data["chapter"] = inferred_chapter
data["_audit_file"] = str(audit_file)
audits.append(data)
for issue in data.get("issues", []) or []:
issue_record = dict(issue)
issue_record["chapter"] = data.get("chapter")
issue_record["_audit_file"] = str(audit_file)
issues.append(issue_record)
summary = {
"run_dir": str(run_dir),
"generated_at": now_iso(),
"audit_file_count": len(audits),
"issue_count": len(issues),
"clear_count": len([a for a in audits if a.get("status") == "clear"]),
"issues_found_count": len([a for a in audits if a.get("status") == "issues_found"]),
"issues": issues,
}
write_yaml(run_dir / "summary.yml", summary)
print(f"Audit files: {summary['audit_file_count']}")
print(f"Issues: {summary['issue_count']}")
print(f"Summary: {run_dir / 'summary.yml'}")
return 0
def add_selection_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument("--volumes", default="vol1,vol2", help="Comma-separated volumes.")
parser.add_argument("--only", help="Comma-separated units such as vol1:introduction.")
parser.add_argument("--limit", type=int, help="Limit selected units after filtering.")
parser.add_argument("--include-frontmatter", action="store_true")
parser.add_argument("--include-parts", action="store_true")
parser.add_argument("--include-glossary", action="store_true")
parser.add_argument("--include-backmatter", action="store_true")
parser.add_argument("--include-all-units", action="store_true")
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
subparsers = parser.add_subparsers(dest="command", required=True)
list_parser = subparsers.add_parser("list", help="List selected build units.")
add_selection_args(list_parser)
list_parser.set_defaults(func=cmd_list)
build_parser_ = subparsers.add_parser("build", help="Build packets and prompts.")
add_selection_args(build_parser_)
build_parser_.add_argument("--mode", choices=["chapters", "volumes", "both"], default="chapters")
build_parser_.add_argument("--run-id", help="Run directory name under review/math-audit/runs.")
build_parser_.add_argument("--batch-size", type=int, default=3)
build_parser_.add_argument("--force", action="store_true", help="Replace existing packet dirs.")
build_parser_.add_argument("--continue-on-error", action="store_true")
build_parser_.add_argument("--verbose", action="store_true", help="Stream binder output.")
build_parser_.set_defaults(func=cmd_build)
prompts_parser = subparsers.add_parser("prompts", help="Regenerate agent batch prompts.")
prompts_parser.add_argument("run_dir")
prompts_parser.add_argument("--batch-size", type=int, default=3)
prompts_parser.set_defaults(func=cmd_prompts)
summary_parser = subparsers.add_parser("summarize", help="Summarize agent YAML audits.")
summary_parser.add_argument("run_dir")
summary_parser.set_defaults(func=cmd_summarize)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
return args.func(args)
if __name__ == "__main__":
raise SystemExit(main())