From 7ff92d5d44bda8964d726cd6633e50a4b44602c5 Mon Sep 17 00:00:00 2001 From: Joe Harris Date: Mon, 12 Oct 2026 10:17:24 -0400 Subject: [PATCH] refactor(scripts): consolidate the Polars version-matrix fork onto the engine-neutral core Extend MatrixSpec with the setup dimension, validation hooks, and run controls so the Polars matrix is a per-engine spec with thin wrappers, matching the DuckDB shape. The shared runner and analyzer now serve both engines; the Polars scripts delegate to them with identical dry-run output and analysis results. --- scripts/analyze_polars_version_matrix.py | 304 ++------- scripts/analyze_version_matrix.py | 270 +++++++- scripts/polars_matrix_spec.py | 128 +--- scripts/run_polars_version_matrix.py | 644 ++---------------- scripts/run_version_matrix.py | 558 ++++++++++++++- scripts/version_matrix_specs.py | 131 +++- .../scripts/test_polars_version_matrix.py | 73 ++ 7 files changed, 1145 insertions(+), 963 deletions(-) diff --git a/scripts/analyze_polars_version_matrix.py b/scripts/analyze_polars_version_matrix.py index ef4602fefd..21888a5bef 100644 --- a/scripts/analyze_polars_version_matrix.py +++ b/scripts/analyze_polars_version_matrix.py @@ -2,283 +2,57 @@ from __future__ import annotations -import argparse -import csv -import hashlib -import json -import statistics import sys from pathlib import Path -from typing import Any REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from scripts.polars_matrix_spec import POLARS, PolarsMatrixSpec +from scripts import analyze_version_matrix +from scripts.analyze_version_matrix import ( + ACCEPTED_STATUSES, + EVIDENCE_REQUIRED, + PINNED_KEYS, + CellKey, + _stats, + _totals, + analyze_polars as analyze, + assess_cells, + assess_record, + compare_cells as compare, + group_records, + load_manifests, + load_qualification, + timing_stats, + write_polars_outputs as write_outputs, +) +from scripts.version_matrix_specs import POLARS + +__all__ = [ + "ACCEPTED_STATUSES", + "EVIDENCE_REQUIRED", + "PINNED_KEYS", + "CellKey", + "REPO_ROOT", + "_stats", + "_totals", + "analyze", + "assess_cells", + "assess_record", + "compare", + "group_records", + "load_manifests", + "load_qualification", + "timing_stats", + "write_outputs", +] CLI_DESCRIPTION = "Analyze Polars version-matrix manifests against the qualification file." -ACCEPTED_STATUSES = {"SUCCESS", "PASS", "PASSED"} -EVIDENCE_REQUIRED = {"uncertain", "not_run"} -PINNED_KEYS = ("benchbox_commit", "pyarrow", "thread_count", "datagen_manifests") - -CellKey = tuple[str, str, str] - - -def load_qualification(path: Path, spec: PolarsMatrixSpec) -> dict[str, Any]: - raw = path.read_bytes() - digest = hashlib.sha256(raw).hexdigest() - if digest != spec.qualification_sha256: - raise ValueError(f"qualification file hash {digest} differs from the spec hash {spec.qualification_sha256}") - payload = json.loads(raw) - clean = { - (row["version"], row["setup"], row["benchmark"]): bool(row["clean"]) for row in payload.get("eligibility", ()) - } - reference = payload.get("reference", {}).get("row_counts", {}) - divergences = payload.get("known_divergences", {}) - return { - "scale_factor": float(payload["scale_factor"]), - "clean": clean, - "reference": reference, - "divergences": divergences, - "sha256": digest, - } - - -def load_manifests(paths: list[Path], spec: PolarsMatrixSpec) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - manifests = [json.loads(path.read_text(encoding="utf-8")) for path in paths] - if not manifests: - raise ValueError("no manifests given") - for path, manifest in zip(paths, manifests, strict=True): - if manifest.get("spec") != spec.snapshot(): - raise ValueError(f"{path} was produced by a different spec") - if not manifest.get("complete"): - raise ValueError(f"{path} is incomplete") - if manifest.get("workload_seed") is not None: - raise ValueError(f"{path} used a workload seed") - if any(manifest["pins"].get(key) != manifests[0]["pins"].get(key) for key in PINNED_KEYS): - raise ValueError(f"{path} has different pinned values from {paths[0]}") - records = [record for manifest in manifests for record in manifest["records"]] - return manifests, records - - -def group_records(records: list[dict[str, Any]], spec: PolarsMatrixSpec) -> dict[CellKey, dict[int, dict[str, Any]]]: - grouped: dict[CellKey, dict[int, dict[str, Any]]] = {} - for record in records: - key = (record["benchmark"], record["version"], record["setup"]) - rounds = grouped.setdefault(key, {}) - if record["round"] in rounds: - raise ValueError(f"duplicate record for {key} round {record['round']}") - rounds[record["round"]] = record - expected_rounds = set(range(1, spec.rounds + 1)) - for cell in spec.cells(): - key = (cell.benchmark, cell.version, cell.setup) - if set(grouped.get(key, {})) != expected_rounds: - raise ValueError(f"cell {cell.cell_id} does not have exactly rounds {sorted(expected_rounds)}") - unexpected = set(grouped) - {(c.benchmark, c.version, c.setup) for c in spec.cells()} - if unexpected: - raise ValueError(f"records for cells outside the spec: {sorted(unexpected)}") - return grouped - - -def assess_record(spec: PolarsMatrixSpec, qualification: dict[str, Any], record: dict[str, Any]) -> list[str]: - benchmark, version, setup = record["benchmark"], record["version"], record["setup"] - if record["outcome"] != "completed": - return [f"outcome {record['outcome']}"] - reasons: list[str] = [] - if float(record["scale"]) != qualification["scale_factor"]: - reasons.append(f"scale {record['scale']} differs from the qualification scale {qualification['scale_factor']}") - if record.get("rechunk_effective") is not False: - reasons.append("rechunk_effective is not false") - receipt = record.get("execution_engine") if isinstance(record.get("execution_engine"), dict) else {} - requested = receipt.get("requested", record.get("engine_requested")) - if requested != record["engine"]: - reasons.append(f"engine_requested {requested} differs from {record['engine']}") - if record.get("polars_version") != version or record.get("installed", {}).get("polars") != version: - reasons.append("recorded Polars version differs from the cell version") - if record.get("polars_runtime_version") not in (version, "absent"): - reasons.append(f"runtime version {record.get('polars_runtime_version')} differs from {version}") - - validation = record.get("validation") - if validation in EVIDENCE_REQUIRED: - if not qualification["clean"].get((version, setup, benchmark), False): - reasons.append(f"validation {validation} without clean qualification evidence") - elif validation != "passed": - reasons.append(f"validation status {validation}") - - queries = record.get("queries", {}) - if any(str(query.get("status", "")).upper() not in ACCEPTED_STATUSES for query in queries.values()): - reasons.append("a measured query did not succeed") - reference = qualification["reference"].get(benchmark, {}) - streams = {key.split(":", 1)[0] for key in queries} - expected_keys = {key for key in reference if key.split(":", 1)[0] in streams} - if set(queries) != expected_keys: - reasons.append("measured query set differs from the reference query set") - divergences = qualification["divergences"].get(benchmark, {}) - differing = [] - for key, query in queries.items(): - accepted = divergences[key]["observed"] if key in divergences else reference.get(key) - if query.get("rows") != accepted: - differing.append(key) - if differing: - reasons.append(f"row counts differ from the DuckDB reference: {', '.join(sorted(differing)[:5])}") - return reasons - - -def _totals(record: dict[str, Any], excluded: set[str]) -> tuple[float, float]: - queries = record["queries"] - full = sum(float(query["ms"]) for query in queries.values()) - common = sum(float(query["ms"]) for key, query in queries.items() if key not in excluded) - return full, common - - -def _stats(values: list[float]) -> dict[str, float]: - return {"median": float(statistics.median(values)), "min": min(values), "max": max(values)} - - -def timing_stats(records: list[dict[str, Any]], excluded: set[str]) -> dict[str, Any] | None: - try: - totals = [_totals(record, excluded) for record in records] - except (KeyError, TypeError, ValueError): - return None - return {"full_ms": _stats([full for full, _ in totals]), "common_ms": _stats([common for _, common in totals])} - - -def assess_cells( - spec: PolarsMatrixSpec, qualification: dict[str, Any], grouped: dict[CellKey, dict[int, dict[str, Any]]] -) -> list[dict[str, Any]]: - cells = [] - for cell in spec.cells(): - key = (cell.benchmark, cell.version, cell.setup) - records = [grouped[key][index] for index in sorted(grouped[key])] - reasons = [ - f"round {record['round']}: {reason}" - for record in records - for reason in assess_record(spec, qualification, record) - ] - excluded = set(qualification["divergences"].get(cell.benchmark, {})) - stats = timing_stats(records, excluded) - completed = all(record["outcome"] == "completed" for record in records) - if stats is None and not reasons: - reasons.append("measured queries have no timing") - qualified = not reasons - cells.append( - { - "benchmark": cell.benchmark, - "version": cell.version, - "setup": cell.setup, - "engine": cell.engine, - "scale": spec.scales[cell.benchmark], - "rounds": len(records), - "outcomes": [record["outcome"] for record in records], - "qualified": qualified, - "reasons": reasons, - "common_query_exclusions": sorted(excluded), - "peak_rss_gib": max(record["peak_rss_gib"] for record in records), - "swap_out_growth_gib": max(record["swap_out_growth_gib"] for record in records), - "headline": stats if qualified else None, - "reported_separately": stats if completed and not qualified else None, - } - ) - return cells - - -def compare(spec: PolarsMatrixSpec, cells: list[dict[str, Any]]) -> list[dict[str, Any]]: - index = {(cell["benchmark"], cell["version"], cell["setup"]): cell for cell in cells} - comparisons = [] - frames = (("oldest", spec.baseline_version), ("reference", spec.reference_version)) - for cell in cells: - if not cell["qualified"]: - continue - for frame, frame_version in frames: - for baseline_setup in dict.fromkeys((cell["setup"], "A")): - base = index.get((cell["benchmark"], frame_version, baseline_setup)) - if base is not None and base["qualified"]: - break - else: - continue - if base is cell: - continue - value = cell["headline"]["common_ms"] - reference = base["headline"]["common_ms"] - change = value["median"] / reference["median"] - 1 - overlapping = not (value["max"] < reference["min"] or value["min"] > reference["max"]) - comparisons.append( - { - "benchmark": cell["benchmark"], - "version": cell["version"], - "setup": cell["setup"], - "baseline_frame": frame, - "baseline_version": frame_version, - "baseline_setup": baseline_setup, - "median_ms": value["median"], - "baseline_median_ms": reference["median"], - "speedup": reference["median"] / value["median"], - "median_change": change, - "ranges_overlap": overlapping, - "candidate_mover": (not overlapping) and abs(change) >= spec.mover_effect, - } - ) - return comparisons - - -def analyze(spec: PolarsMatrixSpec, manifest_paths: list[Path], qualification_path: Path) -> dict[str, Any]: - qualification = load_qualification(qualification_path, spec) - manifests, records = load_manifests(manifest_paths, spec) - cells = assess_cells(spec, qualification, group_records(records, spec)) - comparisons = compare(spec, cells) - return { - "schema_version": "1", - "source_manifests": [path.name for path in manifest_paths], - "shuffle_seeds": [manifest["shuffle_seed"] for manifest in manifests], - "qualification_sha256": qualification["sha256"], - "rounds": spec.rounds, - "cells": cells, - "comparison_count": len(comparisons), - "comparisons": comparisons, - "candidate_movers": [item for item in comparisons if item["candidate_mover"]], - } - - -def write_outputs(result: dict[str, Any], output_dir: Path) -> None: - (output_dir / "polars-version-matrix-analysis.json").write_text( - json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" - ) - fields = ( - "benchmark", - "version", - "setup", - "baseline_frame", - "baseline_version", - "baseline_setup", - "median_ms", - "baseline_median_ms", - "speedup", - "median_change", - "ranges_overlap", - "candidate_mover", - ) - with (output_dir / "polars-version-matrix-comparisons.csv").open("w", encoding="utf-8", newline="") as handle: - writer = csv.DictWriter(handle, fieldnames=fields) - writer.writeheader() - writer.writerows(result["comparisons"]) def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description=CLI_DESCRIPTION) - parser.add_argument("manifests", type=Path, nargs="+", help="matrix-manifest.json files, one per run directory") - parser.add_argument("--qualification", type=Path, required=True, help="Qualification file named by the spec") - parser.add_argument("--output-dir", type=Path, required=True, help="Directory for the analysis JSON and CSV") - args = parser.parse_args(argv) - try: - result = analyze(POLARS, args.manifests, args.qualification) - args.output_dir.mkdir(parents=True, exist_ok=True) - write_outputs(result, args.output_dir) - except (OSError, ValueError, json.JSONDecodeError, KeyError, TypeError) as exc: - parser.error(str(exc)) - qualified = sum(1 for cell in result["cells"] if cell["qualified"]) - print(f"Analyzed {len(result['cells'])} cells ({qualified} qualified); {result['comparison_count']} comparisons") - return 0 + return analyze_version_matrix.main(argv, spec=POLARS, description=CLI_DESCRIPTION) if __name__ == "__main__": diff --git a/scripts/analyze_version_matrix.py b/scripts/analyze_version_matrix.py index cb37e28ab3..723ceaddaf 100644 --- a/scripts/analyze_version_matrix.py +++ b/scripts/analyze_version_matrix.py @@ -20,6 +20,11 @@ from scripts.version_matrix_specs import SPECS, MatrixSpec CLI_DESCRIPTION = "Compute median engine version-matrix metrics from a run manifest." +ACCEPTED_STATUSES = {"SUCCESS", "PASS", "PASSED"} +EVIDENCE_REQUIRED = {"uncertain", "not_run"} +PINNED_KEYS = ("benchbox_commit", "pyarrow", "thread_count", "datagen_manifests") + +CellKey = tuple[str, str, str] def _median(values: list[float]) -> float | None: @@ -162,7 +167,9 @@ def aggregate_payloads( def _write_json(path: Path, payload: Any) -> None: - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) def write_explorer_bundles(spec: MatrixSpec, manifest_path: Path, output_dir: Path) -> list[Path]: @@ -229,6 +236,245 @@ def analyze(spec: MatrixSpec, manifest_path: Path) -> dict[str, Any]: } +def load_qualification(path: Path, spec: MatrixSpec) -> dict[str, Any]: + raw = path.read_bytes() + digest = hashlib.sha256(raw).hexdigest() + if digest != spec.qualification_sha256: + raise ValueError(f"qualification file hash {digest} differs from the spec hash {spec.qualification_sha256}") + payload = json.loads(raw) + clean = { + (row["version"], row["setup"], row["benchmark"]): bool(row["clean"]) for row in payload.get("eligibility", ()) + } + reference = payload.get("reference", {}).get("row_counts", {}) + divergences = payload.get("known_divergences", {}) + return { + "scale_factor": float(payload["scale_factor"]), + "clean": clean, + "reference": reference, + "divergences": divergences, + "sha256": digest, + } + + +def load_manifests(paths: list[Path], spec: MatrixSpec) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + manifests = [json.loads(path.read_text(encoding="utf-8")) for path in paths] + if not manifests: + raise ValueError("no manifests given") + for path, manifest in zip(paths, manifests, strict=True): + if manifest.get("spec") != spec.snapshot(): + raise ValueError(f"{path} was produced by a different spec") + if not manifest.get("complete"): + raise ValueError(f"{path} is incomplete") + if manifest.get("workload_seed") is not None: + raise ValueError(f"{path} used a workload seed") + if any(manifest["pins"].get(key) != manifests[0]["pins"].get(key) for key in PINNED_KEYS): + raise ValueError(f"{path} has different pinned values from {paths[0]}") + records = [record for manifest in manifests for record in manifest["records"]] + return manifests, records + + +def group_records(records: list[dict[str, Any]], spec: MatrixSpec) -> dict[CellKey, dict[int, dict[str, Any]]]: + grouped: dict[CellKey, dict[int, dict[str, Any]]] = {} + for record in records: + key = (record["benchmark"], record["version"], record["setup"]) + rounds = grouped.setdefault(key, {}) + if record["round"] in rounds: + raise ValueError(f"duplicate record for {key} round {record['round']}") + rounds[record["round"]] = record + expected_rounds = set(range(1, spec.effective_rounds + 1)) + for cell in spec.cells(): + key = (cell.benchmark, cell.version, cell.setup) + if set(grouped.get(key, {})) != expected_rounds: + raise ValueError(f"cell {cell.cell_id} does not have exactly rounds {sorted(expected_rounds)}") + unexpected = set(grouped) - {(c.benchmark, c.version, c.setup) for c in spec.cells()} + if unexpected: + raise ValueError(f"records for cells outside the spec: {sorted(unexpected)}") + return grouped + + +def assess_record(spec: MatrixSpec, qualification: dict[str, Any], record: dict[str, Any]) -> list[str]: + benchmark, version, setup = record["benchmark"], record["version"], record["setup"] + if record["outcome"] != "completed": + return [f"outcome {record['outcome']}"] + reasons: list[str] = [] + if float(record["scale"]) != qualification["scale_factor"]: + reasons.append(f"scale {record['scale']} differs from the qualification scale {qualification['scale_factor']}") + if record.get("rechunk_effective") is not False: + reasons.append("rechunk_effective is not false") + receipt = record.get("execution_engine") if isinstance(record.get("execution_engine"), dict) else {} + requested = receipt.get("requested", record.get("engine_requested")) + if requested != record["engine"]: + reasons.append(f"engine_requested {requested} differs from {record['engine']}") + if record.get("polars_version") != version or record.get("installed", {}).get("polars") != version: + reasons.append("recorded Polars version differs from the cell version") + if record.get("polars_runtime_version") not in (version, "absent"): + reasons.append(f"runtime version {record.get('polars_runtime_version')} differs from {version}") + + validation = record.get("validation") + if validation in EVIDENCE_REQUIRED: + if not qualification["clean"].get((version, setup, benchmark), False): + reasons.append(f"validation {validation} without clean qualification evidence") + elif validation != "passed": + reasons.append(f"validation status {validation}") + + queries = record.get("queries", {}) + if any(str(query.get("status", "")).upper() not in ACCEPTED_STATUSES for query in queries.values()): + reasons.append("a measured query did not succeed") + reference = qualification["reference"].get(benchmark, {}) + streams = {key.split(":", 1)[0] for key in queries} + expected_keys = {key for key in reference if key.split(":", 1)[0] in streams} + if set(queries) != expected_keys: + reasons.append("measured query set differs from the reference query set") + divergences = qualification["divergences"].get(benchmark, {}) + differing = [] + for key, query in queries.items(): + accepted = divergences[key]["observed"] if key in divergences else reference.get(key) + if query.get("rows") != accepted: + differing.append(key) + if differing: + reasons.append(f"row counts differ from the DuckDB reference: {', '.join(sorted(differing)[:5])}") + return reasons + + +def _totals(record: dict[str, Any], excluded: set[str]) -> tuple[float, float]: + queries = record["queries"] + full = sum(float(query["ms"]) for query in queries.values()) + common = sum(float(query["ms"]) for key, query in queries.items() if key not in excluded) + return full, common + + +def _stats(values: list[float]) -> dict[str, float]: + return {"median": float(statistics.median(values)), "min": min(values), "max": max(values)} + + +def timing_stats(records: list[dict[str, Any]], excluded: set[str]) -> dict[str, Any] | None: + try: + totals = [_totals(record, excluded) for record in records] + except (KeyError, TypeError, ValueError): + return None + return {"full_ms": _stats([full for full, _ in totals]), "common_ms": _stats([common for _, common in totals])} + + +def assess_cells( + spec: MatrixSpec, qualification: dict[str, Any], grouped: dict[CellKey, dict[int, dict[str, Any]]] +) -> list[dict[str, Any]]: + cells = [] + for cell in spec.cells(): + key = (cell.benchmark, cell.version, cell.setup) + records = [grouped[key][index] for index in sorted(grouped[key])] + reasons = [ + f"round {record['round']}: {reason}" + for record in records + for reason in assess_record(spec, qualification, record) + ] + excluded = set(qualification["divergences"].get(cell.benchmark, {})) + stats = timing_stats(records, excluded) + completed = all(record["outcome"] == "completed" for record in records) + if stats is None and not reasons: + reasons.append("measured queries have no timing") + qualified = not reasons + cells.append( + { + "benchmark": cell.benchmark, + "version": cell.version, + "setup": cell.setup, + "engine": cell.engine, + "scale": spec.scales[cell.benchmark], + "rounds": len(records), + "outcomes": [record["outcome"] for record in records], + "qualified": qualified, + "reasons": reasons, + "common_query_exclusions": sorted(excluded), + "peak_rss_gib": max(record["peak_rss_gib"] for record in records), + "swap_out_growth_gib": max(record["swap_out_growth_gib"] for record in records), + "headline": stats if qualified else None, + "reported_separately": stats if completed and not qualified else None, + } + ) + return cells + + +def compare_cells(spec: MatrixSpec, cells: list[dict[str, Any]]) -> list[dict[str, Any]]: + index = {(cell["benchmark"], cell["version"], cell["setup"]): cell for cell in cells} + comparisons = [] + frames = (("oldest", spec.baseline_version), ("reference", spec.reference_version)) + for cell in cells: + if not cell["qualified"]: + continue + for frame, frame_version in frames: + for baseline_setup in dict.fromkeys((cell["setup"], "A")): + base = index.get((cell["benchmark"], frame_version, baseline_setup)) + if base is not None and base["qualified"]: + break + else: + continue + if base is cell: + continue + value = cell["headline"]["common_ms"] + reference = base["headline"]["common_ms"] + change = value["median"] / reference["median"] - 1 + overlapping = not (value["max"] < reference["min"] or value["min"] > reference["max"]) + comparisons.append( + { + "benchmark": cell["benchmark"], + "version": cell["version"], + "setup": cell["setup"], + "baseline_frame": frame, + "baseline_version": frame_version, + "baseline_setup": baseline_setup, + "median_ms": value["median"], + "baseline_median_ms": reference["median"], + "speedup": reference["median"] / value["median"], + "median_change": change, + "ranges_overlap": overlapping, + "candidate_mover": (not overlapping) and abs(change) >= spec.mover_effect, + } + ) + return comparisons + + +def analyze_polars(spec: MatrixSpec, manifest_paths: list[Path], qualification_path: Path) -> dict[str, Any]: + qualification = load_qualification(qualification_path, spec) + manifests, records = load_manifests(manifest_paths, spec) + cells = assess_cells(spec, qualification, group_records(records, spec)) + comparisons = compare_cells(spec, cells) + return { + "schema_version": "1", + "source_manifests": [path.name for path in manifest_paths], + "shuffle_seeds": [manifest["shuffle_seed"] for manifest in manifests], + "qualification_sha256": qualification["sha256"], + "rounds": spec.effective_rounds, + "cells": cells, + "comparison_count": len(comparisons), + "comparisons": comparisons, + "candidate_movers": [item for item in comparisons if item["candidate_mover"]], + } + + +def write_polars_outputs(result: dict[str, Any], output_dir: Path) -> None: + (output_dir / "polars-version-matrix-analysis.json").write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + fields = ( + "benchmark", + "version", + "setup", + "baseline_frame", + "baseline_version", + "baseline_setup", + "median_ms", + "baseline_median_ms", + "speedup", + "median_change", + "ranges_overlap", + "candidate_mover", + ) + with (output_dir / "polars-version-matrix-comparisons.csv").open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerows(result["comparisons"]) + + def write_outputs(spec: MatrixSpec, result: dict[str, Any], output_dir: Path) -> None: (output_dir / f"{spec.analysis_stem}-analysis.json").write_text( json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" @@ -253,6 +499,7 @@ def main(argv: list[str] | None = None, *, spec: MatrixSpec | None = None, descr if spec is None: parser.add_argument("--engine", choices=sorted(SPECS), required=True, help="Engine spec of the manifest") parser.add_argument("manifest", type=Path, help="matrix-manifest.json emitted by the runner") + parser.add_argument("extra_manifests", type=Path, nargs="*", help="additional manifests for a setup matrix") parser.add_argument( "--output-dir", type=Path, @@ -263,8 +510,29 @@ def main(argv: list[str] | None = None, *, spec: MatrixSpec | None = None, descr type=Path, help="also write one median bundle per version/benchmark cell to this directory", ) + parser.add_argument("--qualification", type=Path, default=None, help="Qualification file named by the spec") args = parser.parse_args(argv) selected = spec or SPECS[args.engine] + if selected.setup_versions: + if args.explorer_bundles_dir is not None: + parser.error("--explorer-bundles-dir is not supported for a setup matrix") + if args.qualification is None: + parser.error("--qualification is required for a setup matrix") + if args.output_dir is None: + parser.error("--output-dir is required for a setup matrix") + try: + result = analyze_polars(selected, [args.manifest, *args.extra_manifests], args.qualification) + args.output_dir.mkdir(parents=True, exist_ok=True) + write_polars_outputs(result, args.output_dir) + except (OSError, ValueError, json.JSONDecodeError, KeyError, TypeError) as exc: + parser.error(str(exc)) + qualified = sum(1 for cell in result["cells"] if cell["qualified"]) + print( + f"Analyzed {len(result['cells'])} cells ({qualified} qualified); {result['comparison_count']} comparisons" + ) + return 0 + if args.extra_manifests or args.qualification is not None: + parser.error("extra manifests and --qualification apply only to a setup matrix") try: result = analyze(selected, args.manifest) output_dir = args.output_dir or args.manifest.parent diff --git a/scripts/polars_matrix_spec.py b/scripts/polars_matrix_spec.py index 141e6fddec..fc5408497f 100644 --- a/scripts/polars_matrix_spec.py +++ b/scripts/polars_matrix_spec.py @@ -1,122 +1,12 @@ from __future__ import annotations -import random -from dataclasses import dataclass - -SETUP_ENGINES = {"A": "default", "B": "in-memory", "C": "streaming"} -GIB = 1024**3 - - -@dataclass(frozen=True) -class Cell: - benchmark: str - version: str - setup: str - - @property - def engine(self) -> str: - return SETUP_ENGINES[self.setup] - - @property - def cell_id(self) -> str: - return f"{self.benchmark}/{self.version}/{self.setup}" - - -@dataclass(frozen=True) -class PolarsMatrixSpec: - versions: tuple[str, ...] - setup_versions: tuple[tuple[str, tuple[str, ...]], ...] - benchmarks: tuple[tuple[str, float], ...] - expected_queries: tuple[tuple[str, int], ...] - rounds: int - iterations: int - thread_count: int - cell_timeout_s: float - rss_limit_gib: float - swap_limit_gib: float - safety_swap_gib: float - reference_version: str - mover_effect: float - qualification_sha256: str - platform: str = "polars-df" - - @property - def baseline_version(self) -> str: - return self.versions[0] - - @property - def benchmark_ids(self) -> tuple[str, ...]: - return tuple(benchmark for benchmark, _ in self.benchmarks) - - @property - def scales(self) -> dict[str, float]: - return dict(self.benchmarks) - - @property - def query_counts(self) -> dict[str, int]: - return dict(self.expected_queries) - - @property - def setups(self) -> dict[str, tuple[str, ...]]: - return dict(self.setup_versions) - - def cells(self) -> list[Cell]: - return [ - Cell(benchmark, version, setup) - for benchmark in self.benchmark_ids - for setup, versions in self.setup_versions - for version in versions - ] - - @property - def expected_invocations(self) -> int: - per_benchmark = sum(len(versions) for _, versions in self.setup_versions) - return self.rounds * per_benchmark * len(self.benchmarks) - - def snapshot(self) -> dict[str, object]: - return { - "platform": self.platform, - "versions": list(self.versions), - "setups": {setup: list(versions) for setup, versions in self.setup_versions}, - "engines": dict(SETUP_ENGINES), - "benchmarks": [{"id": name, "scale": scale} for name, scale in self.benchmarks], - "expected_queries": dict(self.expected_queries), - "rounds": self.rounds, - "iterations": self.iterations, - "thread_count": self.thread_count, - "cell_timeout_s": self.cell_timeout_s, - "rss_limit_gib": self.rss_limit_gib, - "swap_limit_gib": self.swap_limit_gib, - "safety_swap_gib": self.safety_swap_gib, - "reference_version": self.reference_version, - "mover_effect": self.mover_effect, - "qualification_sha256": self.qualification_sha256, - } - - -def round_order(spec: PolarsMatrixSpec, shuffle_seed: int, round_index: int) -> list[Cell]: - cells = spec.cells() - random.Random(f"{shuffle_seed}:{round_index}").shuffle(cells) - return cells - - -POLARS = PolarsMatrixSpec( - versions=("1.31.0", "1.35.2", "1.40.1", "1.44.2", "2.0.0"), - setup_versions=( - ("A", ("1.31.0", "1.35.2", "1.40.1", "1.44.2", "2.0.0")), - ("B", ("2.0.0",)), - ("C", ("1.31.0", "1.40.1", "1.44.2", "2.0.0")), - ), - benchmarks=(("tpch", 10.0), ("tpcds", 10.0), ("clickbench", 10.0), ("ssb", 10.0)), - expected_queries=(("tpch", 22), ("tpcds", 103), ("clickbench", 43), ("ssb", 13)), - rounds=3, - iterations=1, - thread_count=10, - cell_timeout_s=1800.0, - rss_limit_gib=14.0, - swap_limit_gib=1.0, - safety_swap_gib=6.0, - reference_version="1.44.2", - mover_effect=0.10, - qualification_sha256="da8ea047afdbc880177f6a295dc8d0e9afd1b95a3f044d906e9c6a26b43d2580", +from scripts.version_matrix_specs import ( + GIB, + POLARS, + SETUP_ENGINES, + Cell, + MatrixSpec as PolarsMatrixSpec, + round_order, ) + +__all__ = ["POLARS", "SETUP_ENGINES", "Cell", "GIB", "PolarsMatrixSpec", "round_order"] diff --git a/scripts/run_polars_version_matrix.py b/scripts/run_polars_version_matrix.py index 29f0467ff3..29051017d1 100644 --- a/scripts/run_polars_version_matrix.py +++ b/scripts/run_polars_version_matrix.py @@ -2,32 +2,81 @@ from __future__ import annotations -import argparse -import hashlib -import json -import os -import platform -import shlex -import shutil -import signal -import subprocess import sys -import threading -import time -from datetime import datetime, timezone -from importlib import metadata from pathlib import Path -from typing import Any - -import psutil REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from benchbox.utils.clock import elapsed_seconds, mono_time -from scripts.polars_matrix_spec import GIB, POLARS, Cell, PolarsMatrixSpec, round_order -from scripts.version_matrix_specs import scale_argument +from scripts import run_version_matrix +from scripts.run_version_matrix import ( + CRITICAL_PRESSURE, + CRITICAL_SAMPLES, + OOM_EXIT_CODES, + RUNTIME_DISTRIBUTIONS, + Cell, + _append_log, + _find_optional_result_path as _find_result_path, + _git, + _kill_group, + _run_logged, + cell_command, + cell_options, + check_pins, + classify_outcome, + current_pins, + directory_bytes, + extract_result, + generate_command, + host_description, + host_state, + monitor, + parse_freeze, + plan_lines, + prepare_cell_environment, + pressure_level, + prime_command, + run_cell, + run_polars_matrix as run_matrix, + swap_out_bytes, + verify_installation, +) +from scripts.version_matrix_specs import POLARS + +__all__ = [ + "CRITICAL_PRESSURE", + "CRITICAL_SAMPLES", + "OOM_EXIT_CODES", + "RUNTIME_DISTRIBUTIONS", + "Cell", + "REPO_ROOT", + "_append_log", + "_find_result_path", + "_git", + "_kill_group", + "_run_logged", + "check_pins", + "cell_command", + "cell_options", + "classify_outcome", + "current_pins", + "directory_bytes", + "extract_result", + "generate_command", + "host_description", + "host_state", + "monitor", + "parse_freeze", + "plan_lines", + "prepare_cell_environment", + "pressure_level", + "prime_command", + "run_cell", + "run_matrix", + "swap_out_bytes", + "verify_installation", +] CLI_DESCRIPTION = ( "Run the Polars version matrix.\n" @@ -43,564 +92,9 @@ "``uv run --no-sync`` so the installed Polars is not reverted between cells.\n" ) -RUNTIME_DISTRIBUTIONS = ("polars-runtime-32", "polars-runtime-64", "polars-runtime-compat") -OOM_EXIT_CODES = (-9, 137) -CRITICAL_PRESSURE = 4 -CRITICAL_SAMPLES = 3 - - -def parse_freeze(text: str) -> dict[str, str]: - packages: dict[str, str] = {} - for line in text.splitlines(): - name, separator, version = line.strip().partition("==") - if separator: - packages[name.lower().replace("_", "-")] = version - return packages - - -def verify_installation(packages: dict[str, str], version: str) -> dict[str, Any]: - if packages.get("polars") != version: - raise RuntimeError(f"expected polars=={version}, found {packages.get('polars')}") - runtimes = {name: value for name, value in packages.items() if name.startswith("polars-runtime-")} - stale = {name: value for name, value in runtimes.items() if value != version} - if stale: - raise RuntimeError(f"stale Polars runtime distributions for {version}: {stale}") - return {"polars": version, "runtimes": runtimes, "pyarrow": packages.get("pyarrow")} - - -def check_pins(expected: dict[str, Any], current: dict[str, Any]) -> None: - changed = sorted(key for key in expected if expected[key] != current.get(key)) - if changed: - raise RuntimeError(f"pinned value changed mid-matrix: {', '.join(changed)}") - - -def classify_outcome( - spec: PolarsMatrixSpec, - *, - exit_code: int, - timed_out: bool, - safety_abort: bool, - result_ok: bool, - peak_rss_gib: float, - swap_growth_gib: float, -) -> str: - if safety_abort: - return "resource-policy-exceeded" - if timed_out: - return "timeout" - if exit_code in OOM_EXIT_CODES: - return "oom-killed" - if exit_code != 0 or not result_ok: - return "execution-failure" - if peak_rss_gib > spec.rss_limit_gib or swap_growth_gib > spec.swap_limit_gib: - return "resource-policy-exceeded" - return "completed" - - -def generate_command(benchmark: str, scale: float) -> list[str]: - return [ - "uv", - "run", - "--no-sync", - "--", - "benchbox", - "run", - "--non-interactive", - "--quiet", - "--phases", - "generate", - "--benchmark", - benchmark, - "--scale", - scale_argument(scale), - ] - - -def cell_options(spec: PolarsMatrixSpec, benchmark: str, scale: float, engine: str) -> list[str]: - return [ - "--platform", - spec.platform, - "--benchmark", - benchmark, - "--scale", - scale_argument(scale), - "--tuning", - "notuning", - "--ignore-memory-warnings", - "--platform-option", - "rechunk=false", - "--execution-engine", - engine, - ] - - -def prime_command(spec: PolarsMatrixSpec, benchmark: str, scale: float) -> list[str]: - base = ["uv", "run", "--no-sync", "--", "benchbox", "run", "--non-interactive", "--quiet", "--phases", "load"] - return [*base, *cell_options(spec, benchmark, scale, "default")] - - -def cell_command(spec: PolarsMatrixSpec, cell: Cell) -> list[str]: - base = [ - "uv", - "run", - "--no-sync", - "--", - "benchbox", - "run", - "--non-interactive", - "--quiet", - "--iterations", - str(spec.iterations), - ] - return [*base, *cell_options(spec, cell.benchmark, spec.scales[cell.benchmark], cell.engine)] - - -def plan_lines(spec: PolarsMatrixSpec, *, shuffle_seed: int, first_round: int, rounds: int) -> list[str]: - lines = [ - f"# {rounds} rounds x {len(spec.cells())} cells = {rounds * len(spec.cells())} invocations " - f"(shuffle seed {shuffle_seed}, first round {first_round})" - ] - lines.extend(f"$ {shlex.join(generate_command(name, scale))}" for name, scale in spec.benchmarks) - lines.extend(f"$ {shlex.join(prime_command(spec, name, scale))}" for name, scale in spec.benchmarks) - for round_index in range(first_round, first_round + rounds): - order = round_order(spec, shuffle_seed, round_index) - for position, cell in enumerate(order, start=1): - lines.append( - f"round {round_index} position {position}/{len(order)} {cell.cell_id}: install polars=={cell.version}; " - f"$ {shlex.join(cell_command(spec, cell))}" - ) - return lines - - -def _append_log(log_path: Path, text: str) -> None: - with log_path.open("a", encoding="utf-8") as handle: - handle.write(text) - if not text.endswith("\n"): - handle.write("\n") - - -def _run_logged(command: list[str], *, cwd: Path, env: dict[str, str], log_path: Path) -> tuple[int, str, float]: - _append_log(log_path, f"$ {shlex.join(command)}\n") - started = mono_time() - completed = subprocess.run(command, cwd=cwd, env=env, capture_output=True, text=True, check=False) - elapsed = elapsed_seconds(started) - output = (completed.stdout or "") + (completed.stderr or "") - _append_log(log_path, output) - _append_log(log_path, f"[exit={completed.returncode} elapsed_s={elapsed:.3f}]\n") - return completed.returncode, completed.stdout or "", elapsed - - -def _find_result_path(output: str, *, cwd: Path) -> Path | None: - for line in reversed(output.splitlines()): - candidate = Path(line.strip()) - if not candidate.is_absolute(): - candidate = cwd / candidate - if candidate.is_file() and candidate.suffix == ".json": - return candidate.resolve() - return None - - -def swap_out_bytes() -> int: - page = os.sysconf("SC_PAGE_SIZE") - text = subprocess.run(["vm_stat"], capture_output=True, text=True, check=False).stdout - for line in text.splitlines(): - if line.startswith("Swapouts"): - return int(line.split(":")[1].strip().rstrip(".")) * page - return 0 - - -def pressure_level() -> int: - text = subprocess.run( - ["sysctl", "-n", "kern.memorystatus_vm_pressure_level"], capture_output=True, text=True, check=False - ).stdout.strip() - return int(text) if text else 1 - - -def directory_bytes(path: Path) -> int: - total = 0 - for item in path.rglob("*"): - try: - if item.is_file(): - total += item.stat().st_size - except OSError: - continue - return total - - -def host_state(path: Path) -> dict[str, Any]: - load = os.getloadavg() - return { - "load_1m": round(load[0], 2), - "pressure_level": pressure_level(), - "free_disk_gib": round(shutil.disk_usage(path).free / GIB, 1), - "swap_out_bytes": swap_out_bytes(), - } - - -def host_description() -> dict[str, Any]: - return { - "platform": platform.platform(), - "machine": platform.machine(), - "cpu_count": os.cpu_count(), - "memory_gib": round(psutil.virtual_memory().total / GIB, 1), - "python": platform.python_version(), - } - - -def _kill_group(process: subprocess.Popen[str]) -> None: - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - - -def monitor( - spec: PolarsMatrixSpec, process: subprocess.Popen[str], spill_dir: Path, swap_start: int, pressure_start: int -) -> dict[str, Any]: - root = psutil.Process(process.pid) - started = mono_time() - peak_rss = 0 - peak_spill = 0 - max_pressure = pressure_start - critical = 0 - samples = 0 - timed_out = False - safety_abort = False - while process.poll() is None: - try: - total = 0 - for member in [root, *root.children(recursive=True)]: - try: - total += member.memory_info().rss - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - peak_rss = max(peak_rss, total) - except psutil.NoSuchProcess: - pass - if samples % 5 == 0: - peak_spill = max(peak_spill, directory_bytes(spill_dir)) - level = pressure_level() - max_pressure = max(max_pressure, level) - critical = critical + 1 if level >= CRITICAL_PRESSURE else 0 - if critical >= CRITICAL_SAMPLES or swap_out_bytes() - swap_start > spec.safety_swap_gib * GIB: - safety_abort = True - _kill_group(process) - break - samples += 1 - if elapsed_seconds(started) > spec.cell_timeout_s: - timed_out = True - _kill_group(process) - break - time.sleep(1) - process.wait() - return { - "peak_rss_bytes": peak_rss, - "peak_spill_bytes": max(peak_spill, directory_bytes(spill_dir)), - "max_pressure": max_pressure, - "timed_out": timed_out, - "safety_abort": safety_abort, - } - - -def extract_result(payload: dict[str, Any], spec: PolarsMatrixSpec, cell: Cell) -> dict[str, Any]: - platform_info = payload.get("platform") if isinstance(payload.get("platform"), dict) else {} - config = platform_info.get("config") if isinstance(platform_info.get("config"), dict) else {} - summary = payload.get("summary") if isinstance(payload.get("summary"), dict) else {} - counts = summary.get("queries") if isinstance(summary.get("queries"), dict) else {} - queries: dict[str, dict[str, Any]] = {} - for query in payload.get("queries", []): - if isinstance(query, dict) and query.get("run_type") == "measurement": - queries[f"{query.get('stream')}:{query.get('id')}"] = { - "rows": query.get("rows"), - "status": query.get("status"), - "ms": query.get("ms"), - } - receipt = platform_info.get("execution_engine") - if not isinstance(receipt, dict): - receipt = config.get("execution_engine") if isinstance(config.get("execution_engine"), dict) else {} - requested = receipt.get("requested", config.get("engine_requested")) - expected = spec.query_counts[cell.benchmark] - checks = { - "rechunk_effective_false": config.get("rechunk_effective") is False, - "engine_matches": requested == cell.engine, - "polars_version_matches": platform_info.get("version") == cell.version, - "query_count_matches": counts.get("total") == expected and len(queries) == expected, - "no_failed_queries": counts.get("failed") == 0, - } - return { - "validation": summary.get("validation"), - "query_total": counts.get("total"), - "query_failed": counts.get("failed"), - "rechunk_effective": config.get("rechunk_effective"), - "execution_engine": receipt, - "engine_requested": requested, - "collect_engine_argument": (receipt.get("applied_native") or {}) - .get("collect_kwargs", {}) - .get("engine", config.get("collect_engine_argument")), - "polars_version": platform_info.get("version"), - "polars_runtime_package": config.get("polars_runtime_package"), - "polars_runtime_version": config.get("polars_runtime_version"), - "queries": queries, - "checks": checks, - "result_ok": all(checks.values()), - } - - -def _git(repo_root: Path, *arguments: str) -> str: - return subprocess.run(["git", *arguments], cwd=repo_root, capture_output=True, text=True, check=True).stdout.strip() - - -def current_pins(repo_root: Path, datagen_root: Path, spec: PolarsMatrixSpec) -> dict[str, Any]: - manifests = {} - for name, _ in spec.benchmarks: - for path in sorted(datagen_root.glob(f"{name}_sf*/_datagen_manifest.json")): - manifests[path.parent.name] = hashlib.sha256(path.read_bytes()).hexdigest() - return { - "benchbox_commit": _git(repo_root, "rev-parse", "HEAD"), - "pyarrow": metadata.version("pyarrow"), - "thread_count": spec.thread_count, - "datagen_manifests": manifests, - } - - -def _write_json(path: Path, payload: Any) -> None: - temporary = path.with_suffix(path.suffix + ".tmp") - temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") - temporary.replace(path) - - -def prepare_cell_environment( - cell: Cell, spec: PolarsMatrixSpec, *, repo_root: Path, env: dict[str, str], log_path: Path -) -> dict[str, Any]: - _run_logged( - ["uv", "pip", "uninstall", "--python", sys.executable, "polars", *RUNTIME_DISTRIBUTIONS], - cwd=repo_root, - env=env, - log_path=log_path, - ) - status, _, _ = _run_logged( - ["uv", "pip", "install", "--python", sys.executable, f"polars=={cell.version}"], - cwd=repo_root, - env=env, - log_path=log_path, - ) - if status: - raise RuntimeError(f"could not install polars=={cell.version}; see {log_path}") - status, frozen, _ = _run_logged( - ["uv", "pip", "freeze", "--python", sys.executable], cwd=repo_root, env=env, log_path=log_path - ) - if status: - raise RuntimeError(f"could not list installed packages; see {log_path}") - installed = verify_installation(parse_freeze(frozen), cell.version) - probe = [sys.executable, "-c", "import polars as pl;print(pl.__version__);print(pl.thread_pool_size())"] - status, probed, _ = _run_logged(probe, cwd=repo_root, env=env, log_path=log_path) - reported = probed.split() - if status or reported != [cell.version, str(spec.thread_count)]: - raise RuntimeError(f"Polars reports {reported}, expected {[cell.version, str(spec.thread_count)]}") - installed["thread_pool_size"] = spec.thread_count - return installed - - -def run_cell( - spec: PolarsMatrixSpec, - cell: Cell, - *, - round_index: int, - position: int, - output_dir: Path, - repo_root: Path, - env: dict[str, str], - log_path: Path, -) -> dict[str, Any]: - installed = prepare_cell_environment(cell, spec, repo_root=repo_root, env=env, log_path=log_path) - spill_dir = output_dir / "spill" / f"round{round_index}-{cell.benchmark}-{cell.version}-{cell.setup}" - if spill_dir.exists(): - shutil.rmtree(spill_dir) - spill_dir.mkdir(parents=True) - cell_env = {**env, "POLARS_TEMP_DIR": str(spill_dir), "POLARS_MAX_THREADS": str(spec.thread_count)} - command = cell_command(spec, cell) - _append_log(log_path, f"$ {shlex.join(command)}\n") - before = host_state(output_dir) - started_at = datetime.now(tz=timezone.utc).isoformat() - started = mono_time() - with log_path.open("a", encoding="utf-8") as handle: - process = subprocess.Popen( - command, - cwd=repo_root, - env=cell_env, - stdout=subprocess.PIPE, - stderr=handle, - text=True, - start_new_session=True, - ) - stdout_lines: list[str] = [] - reader = threading.Thread(target=lambda: stdout_lines.extend(process.stdout or []), daemon=True) - reader.start() - observed = monitor(spec, process, spill_dir, before["swap_out_bytes"], before["pressure_level"]) - reader.join(timeout=30) - elapsed = elapsed_seconds(started) - after = host_state(output_dir) - output = "".join(stdout_lines) - _append_log(log_path, output) - _append_log(log_path, f"[exit={process.returncode} elapsed_s={elapsed:.3f}]\n") - result_path = _find_result_path(output, cwd=repo_root) - extracted: dict[str, Any] = {} - if result_path is not None: - try: - extracted = extract_result(json.loads(result_path.read_text(encoding="utf-8")), spec, cell) - except (OSError, ValueError, TypeError): - extracted = {} - swap_growth_gib = max(0, after["swap_out_bytes"] - before["swap_out_bytes"]) / GIB - peak_rss_gib = observed["peak_rss_bytes"] / GIB - outcome = classify_outcome( - spec, - exit_code=process.returncode, - timed_out=observed["timed_out"], - safety_abort=observed["safety_abort"], - result_ok=bool(extracted.get("result_ok")), - peak_rss_gib=peak_rss_gib, - swap_growth_gib=swap_growth_gib, - ) - shutil.rmtree(spill_dir, ignore_errors=True) - return { - "round": round_index, - "position": position, - "benchmark": cell.benchmark, - "version": cell.version, - "setup": cell.setup, - "engine": cell.engine, - "scale": spec.scales[cell.benchmark], - "started_at": started_at, - "elapsed_s": round(elapsed, 3), - "exit_code": process.returncode, - "outcome": outcome, - "host_safety_abort": observed["safety_abort"], - "peak_rss_gib": round(peak_rss_gib, 3), - "swap_out_growth_gib": round(swap_growth_gib, 3), - "spill_peak_gib": round(observed["peak_spill_bytes"] / GIB, 3), - "pressure_before": before["pressure_level"], - "pressure_after": after["pressure_level"], - "pressure_max": observed["max_pressure"], - "load_1m_before": before["load_1m"], - "free_disk_before_gib": before["free_disk_gib"], - "free_disk_after_gib": after["free_disk_gib"], - "installed": installed, - "result_path": result_path.relative_to(output_dir).as_posix() - if result_path is not None and output_dir in result_path.parents - else None, - **{key: value for key, value in extracted.items() if key != "result_ok"}, - } - - -def run_matrix( - spec: PolarsMatrixSpec, - *, - output_dir: Path, - repo_root: Path, - shuffle_seed: int, - first_round: int, - rounds: int, - qualification: Path, -) -> dict[str, Any]: - if output_dir == repo_root or repo_root in output_dir.parents: - raise ValueError(f"output directory must be outside the checkout: {output_dir}") - if output_dir.exists() and any(output_dir.iterdir()): - raise ValueError(f"output directory must be new or empty: {output_dir}") - if _git(repo_root, "status", "--porcelain", "--untracked-files=no"): - raise ValueError("the checkout has uncommitted changes") - qualification_sha = hashlib.sha256(qualification.read_bytes()).hexdigest() - if qualification_sha != spec.qualification_sha256: - raise ValueError(f"qualification file hash {qualification_sha} differs from the spec") - output_dir.mkdir(parents=True, exist_ok=True) - shutil.copy(qualification, output_dir / "qualification.json") - log_path = output_dir / "matrix.log" - env = os.environ.copy() - env["BENCHBOX_OUTPUT_DIR"] = str(output_dir) - env["PYTHONUNBUFFERED"] = "1" - env["POLARS_MAX_THREADS"] = str(spec.thread_count) - - for name, scale in spec.benchmarks: - status, _, _ = _run_logged(generate_command(name, scale), cwd=repo_root, env=env, log_path=log_path) - if status: - raise RuntimeError(f"data generation failed for {name}; see {log_path}") - first_cell = Cell(spec.benchmark_ids[0], spec.baseline_version, "A") - prepare_cell_environment(first_cell, spec, repo_root=repo_root, env=env, log_path=log_path) - for name, scale in spec.benchmarks: - status, _, _ = _run_logged(prime_command(spec, name, scale), cwd=repo_root, env=env, log_path=log_path) - if status: - raise RuntimeError(f"Parquet conversion failed for {name}; see {log_path}") - - pins = current_pins(repo_root, output_dir / "datagen", spec) - manifest: dict[str, Any] = { - "schema_version": "1", - "created_at": datetime.now(tz=timezone.utc).isoformat(), - "spec": spec.snapshot(), - "shuffle_seed": shuffle_seed, - "workload_seed": None, - "first_round": first_round, - "rounds": rounds, - "pins": pins, - "host": host_description(), - "round_order": {}, - "records": [], - "complete": False, - } - for round_index in range(first_round, first_round + rounds): - order = round_order(spec, shuffle_seed, round_index) - manifest["round_order"][str(round_index)] = [cell.cell_id for cell in order] - for position, cell in enumerate(order, start=1): - record = run_cell( - spec, - cell, - round_index=round_index, - position=position, - output_dir=output_dir, - repo_root=repo_root, - env=env, - log_path=log_path, - ) - manifest["records"].append(record) - _write_json(output_dir / "matrix-manifest.json", manifest) - check_pins(pins, current_pins(repo_root, output_dir / "datagen", spec)) - manifest["complete"] = True - _write_json(output_dir / "matrix-manifest.json", manifest) - return manifest - def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description=CLI_DESCRIPTION, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--output-dir", type=Path, required=True, help="New output root outside the checkout") - parser.add_argument("--shuffle-seed", type=int, required=True, help="Seed for the per-round cell order only") - parser.add_argument("--first-round", type=int, default=1, help="Index of the first round (default: 1)") - parser.add_argument("--rounds", type=int, default=POLARS.rounds, help="Number of rounds to run") - parser.add_argument("--qualification", type=Path, help="Qualification file whose SHA-256 the spec records") - parser.add_argument("--dry-run", action="store_true", help="Print the planned invocations without running them") - args = parser.parse_args(argv) - if args.dry_run: - print( - "\n".join( - plan_lines(POLARS, shuffle_seed=args.shuffle_seed, first_round=args.first_round, rounds=args.rounds) - ) - ) - return 0 - if args.qualification is None: - parser.error("--qualification is required unless --dry-run is given") - try: - manifest = run_matrix( - POLARS, - output_dir=args.output_dir.expanduser().resolve(), - repo_root=REPO_ROOT, - shuffle_seed=args.shuffle_seed, - first_round=args.first_round, - rounds=args.rounds, - qualification=args.qualification.expanduser().resolve(), - ) - except (OSError, RuntimeError, ValueError, subprocess.CalledProcessError) as exc: - print(f"ERROR: {exc}", file=sys.stderr) - return 1 - print(f"Polars matrix complete: {len(manifest['records'])} runs") - print(f"Manifest: {args.output_dir.expanduser().resolve() / 'matrix-manifest.json'}") - return 0 + return run_version_matrix.main(argv, spec=POLARS, description=CLI_DESCRIPTION) if __name__ == "__main__": diff --git a/scripts/run_version_matrix.py b/scripts/run_version_matrix.py index 5e85585e4a..d60579edfb 100644 --- a/scripts/run_version_matrix.py +++ b/scripts/run_version_matrix.py @@ -3,13 +3,20 @@ from __future__ import annotations import argparse +import hashlib import json import os +import platform import shlex +import shutil +import signal import subprocess import sys +import threading +import time from dataclasses import dataclass from datetime import datetime, timezone +from importlib import metadata from pathlib import Path from typing import Any @@ -17,8 +24,15 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) +import psutil + from benchbox.utils.clock import elapsed_seconds, mono_time -from scripts.version_matrix_specs import SPECS, MatrixSpec, scale_argument, scale_token +from scripts.version_matrix_specs import GIB, SPECS, Cell, MatrixSpec, round_order, scale_argument, scale_token + +RUNTIME_DISTRIBUTIONS = ("polars-runtime-32", "polars-runtime-64", "polars-runtime-compat") +OOM_EXIT_CODES = (-9, 137) +CRITICAL_PRESSURE = 4 +CRITICAL_SAMPLES = 3 CLI_DESCRIPTION = ( "Run a reproducible engine version matrix.\n" @@ -200,7 +214,132 @@ def _record( def _write_json(path: Path, payload: Any) -> None: - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def parse_freeze(text: str) -> dict[str, str]: + packages: dict[str, str] = {} + for line in text.splitlines(): + name, separator, version = line.strip().partition("==") + if separator: + packages[name.lower().replace("_", "-")] = version + return packages + + +def verify_installation(packages: dict[str, str], version: str) -> dict[str, Any]: + if packages.get("polars") != version: + raise RuntimeError(f"expected polars=={version}, found {packages.get('polars')}") + runtimes = {name: value for name, value in packages.items() if name.startswith("polars-runtime-")} + stale = {name: value for name, value in runtimes.items() if value != version} + if stale: + raise RuntimeError(f"stale Polars runtime distributions for {version}: {stale}") + return {"polars": version, "runtimes": runtimes, "pyarrow": packages.get("pyarrow")} + + +def check_pins(expected: dict[str, Any], current: dict[str, Any]) -> None: + changed = sorted(key for key in expected if expected[key] != current.get(key)) + if changed: + raise RuntimeError(f"pinned value changed mid-matrix: {', '.join(changed)}") + + +def classify_outcome( + spec: MatrixSpec, + *, + exit_code: int, + timed_out: bool, + safety_abort: bool, + result_ok: bool, + peak_rss_gib: float, + swap_growth_gib: float, +) -> str: + if safety_abort: + return "resource-policy-exceeded" + if timed_out: + return "timeout" + if exit_code in OOM_EXIT_CODES: + return "oom-killed" + if exit_code != 0 or not result_ok: + return "execution-failure" + if peak_rss_gib > spec.rss_limit_gib or swap_growth_gib > spec.swap_limit_gib: + return "resource-policy-exceeded" + return "completed" + + +def generate_command(benchmark: str, scale: float) -> list[str]: + return [ + "uv", + "run", + "--no-sync", + "--", + "benchbox", + "run", + "--non-interactive", + "--quiet", + "--phases", + "generate", + "--benchmark", + benchmark, + "--scale", + scale_argument(scale), + ] + + +def cell_options(spec: MatrixSpec, benchmark: str, scale: float, engine: str) -> list[str]: + return [ + "--platform", + spec.platform, + "--benchmark", + benchmark, + "--scale", + scale_argument(scale), + "--tuning", + "notuning", + "--ignore-memory-warnings", + "--platform-option", + "rechunk=false", + "--execution-engine", + engine, + ] + + +def prime_command(spec: MatrixSpec, benchmark: str, scale: float) -> list[str]: + base = ["uv", "run", "--no-sync", "--", "benchbox", "run", "--non-interactive", "--quiet", "--phases", "load"] + return [*base, *cell_options(spec, benchmark, scale, "default")] + + +def cell_command(spec: MatrixSpec, cell: Cell) -> list[str]: + base = [ + "uv", + "run", + "--no-sync", + "--", + "benchbox", + "run", + "--non-interactive", + "--quiet", + "--iterations", + str(spec.iterations), + ] + return [*base, *cell_options(spec, cell.benchmark, spec.scales[cell.benchmark], cell.engine)] + + +def plan_lines(spec: MatrixSpec, *, shuffle_seed: int, first_round: int, rounds: int) -> list[str]: + lines = [ + f"# {rounds} rounds x {len(spec.cells())} cells = {rounds * len(spec.cells())} invocations " + f"(shuffle seed {shuffle_seed}, first round {first_round})" + ] + lines.extend(f"$ {shlex.join(generate_command(name, scale))}" for name, scale in spec.benchmarks) + lines.extend(f"$ {shlex.join(prime_command(spec, name, scale))}" for name, scale in spec.benchmarks) + for round_index in range(first_round, first_round + rounds): + order = round_order(spec, shuffle_seed, round_index) + for position, cell in enumerate(order, start=1): + lines.append( + f"round {round_index} position {position}/{len(order)} {cell.cell_id}: install polars=={cell.version}; " + f"$ {shlex.join(cell_command(spec, cell))}" + ) + return lines def run_matrix(spec: MatrixSpec, *, output_dir: Path, repo_root: Path, dry_run: bool) -> dict[str, Any]: @@ -270,16 +409,431 @@ def run_matrix(spec: MatrixSpec, *, output_dir: Path, repo_root: Path, dry_run: return manifest +def _run_logged(command: list[str], *, cwd: Path, env: dict[str, str], log_path: Path) -> tuple[int, str, float]: + _append_log(log_path, f"$ {shlex.join(command)}\n") + started = mono_time() + completed = subprocess.run(command, cwd=cwd, env=env, capture_output=True, text=True, check=False) + elapsed = elapsed_seconds(started) + output = (completed.stdout or "") + (completed.stderr or "") + _append_log(log_path, output) + _append_log(log_path, f"[exit={completed.returncode} elapsed_s={elapsed:.3f}]\n") + return completed.returncode, completed.stdout or "", elapsed + + +def _find_optional_result_path(output: str, *, cwd: Path) -> Path | None: + for line in reversed(output.splitlines()): + candidate = Path(line.strip()) + if not candidate.is_absolute(): + candidate = cwd / candidate + if candidate.is_file() and candidate.suffix == ".json": + return candidate.resolve() + return None + + +def swap_out_bytes() -> int: + page = os.sysconf("SC_PAGE_SIZE") + text = subprocess.run(["vm_stat"], capture_output=True, text=True, check=False).stdout + for line in text.splitlines(): + if line.startswith("Swapouts"): + return int(line.split(":")[1].strip().rstrip(".")) * page + return 0 + + +def pressure_level() -> int: + text = subprocess.run( + ["sysctl", "-n", "kern.memorystatus_vm_pressure_level"], capture_output=True, text=True, check=False + ).stdout.strip() + return int(text) if text else 1 + + +def directory_bytes(path: Path) -> int: + total = 0 + for item in path.rglob("*"): + try: + if item.is_file(): + total += item.stat().st_size + except OSError: + continue + return total + + +def host_state(path: Path) -> dict[str, Any]: + load = os.getloadavg() + return { + "load_1m": round(load[0], 2), + "pressure_level": pressure_level(), + "free_disk_gib": round(shutil.disk_usage(path).free / GIB, 1), + "swap_out_bytes": swap_out_bytes(), + } + + +def host_description() -> dict[str, Any]: + return { + "platform": platform.platform(), + "machine": platform.machine(), + "cpu_count": os.cpu_count(), + "memory_gib": round(psutil.virtual_memory().total / GIB, 1), + "python": platform.python_version(), + } + + +def _kill_group(process: subprocess.Popen[str]) -> None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +def monitor( + spec: MatrixSpec, process: subprocess.Popen[str], spill_dir: Path, swap_start: int, pressure_start: int +) -> dict[str, Any]: + root = psutil.Process(process.pid) + started = mono_time() + peak_rss = 0 + peak_spill = 0 + max_pressure = pressure_start + critical = 0 + samples = 0 + timed_out = False + safety_abort = False + while process.poll() is None: + try: + total = 0 + for member in [root, *root.children(recursive=True)]: + try: + total += member.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + peak_rss = max(peak_rss, total) + except psutil.NoSuchProcess: + pass + if samples % 5 == 0: + peak_spill = max(peak_spill, directory_bytes(spill_dir)) + level = pressure_level() + max_pressure = max(max_pressure, level) + critical = critical + 1 if level >= CRITICAL_PRESSURE else 0 + if critical >= CRITICAL_SAMPLES or swap_out_bytes() - swap_start > spec.safety_swap_gib * GIB: + safety_abort = True + _kill_group(process) + break + samples += 1 + if elapsed_seconds(started) > spec.cell_timeout_s: + timed_out = True + _kill_group(process) + break + time.sleep(1) + process.wait() + return { + "peak_rss_bytes": peak_rss, + "peak_spill_bytes": max(peak_spill, directory_bytes(spill_dir)), + "max_pressure": max_pressure, + "timed_out": timed_out, + "safety_abort": safety_abort, + } + + +def extract_result(payload: dict[str, Any], spec: MatrixSpec, cell: Cell) -> dict[str, Any]: + platform_info = payload.get("platform") if isinstance(payload.get("platform"), dict) else {} + config = platform_info.get("config") if isinstance(platform_info.get("config"), dict) else {} + summary = payload.get("summary") if isinstance(payload.get("summary"), dict) else {} + counts = summary.get("queries") if isinstance(summary.get("queries"), dict) else {} + queries: dict[str, dict[str, Any]] = {} + for query in payload.get("queries", []): + if isinstance(query, dict) and query.get("run_type") == "measurement": + queries[f"{query.get('stream')}:{query.get('id')}"] = { + "rows": query.get("rows"), + "status": query.get("status"), + "ms": query.get("ms"), + } + receipt = platform_info.get("execution_engine") + if not isinstance(receipt, dict): + receipt = config.get("execution_engine") if isinstance(config.get("execution_engine"), dict) else {} + requested = receipt.get("requested", config.get("engine_requested")) + expected = spec.query_counts[cell.benchmark] + checks = { + "rechunk_effective_false": config.get("rechunk_effective") is False, + "engine_matches": requested == cell.engine, + "polars_version_matches": platform_info.get("version") == cell.version, + "query_count_matches": counts.get("total") == expected and len(queries) == expected, + "no_failed_queries": counts.get("failed") == 0, + } + return { + "validation": summary.get("validation"), + "query_total": counts.get("total"), + "query_failed": counts.get("failed"), + "rechunk_effective": config.get("rechunk_effective"), + "execution_engine": receipt, + "engine_requested": requested, + "collect_engine_argument": (receipt.get("applied_native") or {}) + .get("collect_kwargs", {}) + .get("engine", config.get("collect_engine_argument")), + "polars_version": platform_info.get("version"), + "polars_runtime_package": config.get("polars_runtime_package"), + "polars_runtime_version": config.get("polars_runtime_version"), + "queries": queries, + "checks": checks, + "result_ok": all(checks.values()), + } + + +def _git(repo_root: Path, *arguments: str) -> str: + return subprocess.run(["git", *arguments], cwd=repo_root, capture_output=True, text=True, check=True).stdout.strip() + + +def current_pins(repo_root: Path, datagen_root: Path, spec: MatrixSpec) -> dict[str, Any]: + manifests = {} + for name, _ in spec.benchmarks: + for path in sorted(datagen_root.glob(f"{name}_sf*/_datagen_manifest.json")): + manifests[path.parent.name] = hashlib.sha256(path.read_bytes()).hexdigest() + return { + "benchbox_commit": _git(repo_root, "rev-parse", "HEAD"), + "pyarrow": metadata.version("pyarrow"), + "thread_count": spec.thread_count, + "datagen_manifests": manifests, + } + + +def prepare_cell_environment( + cell: Cell, spec: MatrixSpec, *, repo_root: Path, env: dict[str, str], log_path: Path +) -> dict[str, Any]: + _run_logged( + ["uv", "pip", "uninstall", "--python", sys.executable, "polars", *RUNTIME_DISTRIBUTIONS], + cwd=repo_root, + env=env, + log_path=log_path, + ) + status, _, _ = _run_logged( + ["uv", "pip", "install", "--python", sys.executable, f"polars=={cell.version}"], + cwd=repo_root, + env=env, + log_path=log_path, + ) + if status: + raise RuntimeError(f"could not install polars=={cell.version}; see {log_path}") + status, frozen, _ = _run_logged( + ["uv", "pip", "freeze", "--python", sys.executable], cwd=repo_root, env=env, log_path=log_path + ) + if status: + raise RuntimeError(f"could not list installed packages; see {log_path}") + installed = verify_installation(parse_freeze(frozen), cell.version) + probe = [sys.executable, "-c", "import polars as pl;print(pl.__version__);print(pl.thread_pool_size())"] + status, probed, _ = _run_logged(probe, cwd=repo_root, env=env, log_path=log_path) + reported = probed.split() + if status or reported != [cell.version, str(spec.thread_count)]: + raise RuntimeError(f"Polars reports {reported}, expected {[cell.version, str(spec.thread_count)]}") + installed["thread_pool_size"] = spec.thread_count + return installed + + +def run_cell( + spec: MatrixSpec, + cell: Cell, + *, + round_index: int, + position: int, + output_dir: Path, + repo_root: Path, + env: dict[str, str], + log_path: Path, +) -> dict[str, Any]: + installed = prepare_cell_environment(cell, spec, repo_root=repo_root, env=env, log_path=log_path) + spill_dir = output_dir / "spill" / f"round{round_index}-{cell.benchmark}-{cell.version}-{cell.setup}" + if spill_dir.exists(): + shutil.rmtree(spill_dir) + spill_dir.mkdir(parents=True) + cell_env = {**env, "POLARS_TEMP_DIR": str(spill_dir), "POLARS_MAX_THREADS": str(spec.thread_count)} + command = cell_command(spec, cell) + _append_log(log_path, f"$ {shlex.join(command)}\n") + before = host_state(output_dir) + started_at = datetime.now(tz=timezone.utc).isoformat() + started = mono_time() + with log_path.open("a", encoding="utf-8") as handle: + process = subprocess.Popen( + command, + cwd=repo_root, + env=cell_env, + stdout=subprocess.PIPE, + stderr=handle, + text=True, + start_new_session=True, + ) + stdout_lines: list[str] = [] + reader = threading.Thread(target=lambda: stdout_lines.extend(process.stdout or []), daemon=True) + reader.start() + observed = monitor(spec, process, spill_dir, before["swap_out_bytes"], before["pressure_level"]) + reader.join(timeout=30) + elapsed = elapsed_seconds(started) + after = host_state(output_dir) + output = "".join(stdout_lines) + _append_log(log_path, output) + _append_log(log_path, f"[exit={process.returncode} elapsed_s={elapsed:.3f}]\n") + result_path = _find_optional_result_path(output, cwd=repo_root) + extracted: dict[str, Any] = {} + if result_path is not None: + try: + extracted = extract_result(json.loads(result_path.read_text(encoding="utf-8")), spec, cell) + except (OSError, ValueError, TypeError): + extracted = {} + swap_growth_gib = max(0, after["swap_out_bytes"] - before["swap_out_bytes"]) / GIB + peak_rss_gib = observed["peak_rss_bytes"] / GIB + outcome = classify_outcome( + spec, + exit_code=process.returncode, + timed_out=observed["timed_out"], + safety_abort=observed["safety_abort"], + result_ok=bool(extracted.get("result_ok")), + peak_rss_gib=peak_rss_gib, + swap_growth_gib=swap_growth_gib, + ) + shutil.rmtree(spill_dir, ignore_errors=True) + return { + "round": round_index, + "position": position, + "benchmark": cell.benchmark, + "version": cell.version, + "setup": cell.setup, + "engine": cell.engine, + "scale": spec.scales[cell.benchmark], + "started_at": started_at, + "elapsed_s": round(elapsed, 3), + "exit_code": process.returncode, + "outcome": outcome, + "host_safety_abort": observed["safety_abort"], + "peak_rss_gib": round(peak_rss_gib, 3), + "swap_out_growth_gib": round(swap_growth_gib, 3), + "spill_peak_gib": round(observed["peak_spill_bytes"] / GIB, 3), + "pressure_before": before["pressure_level"], + "pressure_after": after["pressure_level"], + "pressure_max": observed["max_pressure"], + "load_1m_before": before["load_1m"], + "free_disk_before_gib": before["free_disk_gib"], + "free_disk_after_gib": after["free_disk_gib"], + "installed": installed, + "result_path": result_path.relative_to(output_dir).as_posix() + if result_path is not None and output_dir in result_path.parents + else None, + **{key: value for key, value in extracted.items() if key != "result_ok"}, + } + + +def run_polars_matrix( + spec: MatrixSpec, + *, + output_dir: Path, + repo_root: Path, + shuffle_seed: int, + first_round: int, + rounds: int, + qualification: Path, +) -> dict[str, Any]: + if output_dir == repo_root or repo_root in output_dir.parents: + raise ValueError(f"output directory must be outside the checkout: {output_dir}") + if output_dir.exists() and any(output_dir.iterdir()): + raise ValueError(f"output directory must be new or empty: {output_dir}") + if _git(repo_root, "status", "--porcelain", "--untracked-files=no"): + raise ValueError("the checkout has uncommitted changes") + qualification_sha = hashlib.sha256(qualification.read_bytes()).hexdigest() + if qualification_sha != spec.qualification_sha256: + raise ValueError(f"qualification file hash {qualification_sha} differs from the spec") + output_dir.mkdir(parents=True, exist_ok=True) + shutil.copy(qualification, output_dir / "qualification.json") + log_path = output_dir / "matrix.log" + env = os.environ.copy() + env["BENCHBOX_OUTPUT_DIR"] = str(output_dir) + env["PYTHONUNBUFFERED"] = "1" + env["POLARS_MAX_THREADS"] = str(spec.thread_count) + + for name, scale in spec.benchmarks: + status, _, _ = _run_logged(generate_command(name, scale), cwd=repo_root, env=env, log_path=log_path) + if status: + raise RuntimeError(f"data generation failed for {name}; see {log_path}") + first_cell = Cell(spec.benchmark_ids[0], spec.baseline_version, "A") + prepare_cell_environment(first_cell, spec, repo_root=repo_root, env=env, log_path=log_path) + for name, scale in spec.benchmarks: + status, _, _ = _run_logged(prime_command(spec, name, scale), cwd=repo_root, env=env, log_path=log_path) + if status: + raise RuntimeError(f"Parquet conversion failed for {name}; see {log_path}") + + pins = current_pins(repo_root, output_dir / "datagen", spec) + manifest: dict[str, Any] = { + "schema_version": "1", + "created_at": datetime.now(tz=timezone.utc).isoformat(), + "spec": spec.snapshot(), + "shuffle_seed": shuffle_seed, + "workload_seed": None, + "first_round": first_round, + "rounds": rounds, + "pins": pins, + "host": host_description(), + "round_order": {}, + "records": [], + "complete": False, + } + for round_index in range(first_round, first_round + rounds): + order = round_order(spec, shuffle_seed, round_index) + manifest["round_order"][str(round_index)] = [cell.cell_id for cell in order] + for position, cell in enumerate(order, start=1): + record = run_cell( + spec, + cell, + round_index=round_index, + position=position, + output_dir=output_dir, + repo_root=repo_root, + env=env, + log_path=log_path, + ) + manifest["records"].append(record) + _write_json(output_dir / "matrix-manifest.json", manifest) + check_pins(pins, current_pins(repo_root, output_dir / "datagen", spec)) + manifest["complete"] = True + _write_json(output_dir / "matrix-manifest.json", manifest) + return manifest + + def main(argv: list[str] | None = None, *, spec: MatrixSpec | None = None, description: str = CLI_DESCRIPTION) -> int: parser = argparse.ArgumentParser(description=description) if spec is None: parser.add_argument("--engine", choices=sorted(SPECS), required=True, help="Engine spec to run") parser.add_argument("--output-dir", type=Path, required=True, help="New output root outside the checkout") parser.add_argument("--dry-run", action="store_true", help="Print the planned commands without running them") + parser.add_argument("--shuffle-seed", type=int, default=None, help="Seed for the per-round cell order only") + parser.add_argument("--first-round", type=int, default=1, help="Index of the first round") + parser.add_argument("--rounds", type=int, default=None, help="Number of rounds to run") + parser.add_argument("--qualification", type=Path, default=None, help="Qualification file named by the spec") args = parser.parse_args(argv) selected = spec or SPECS[args.engine] repo_root = REPO_ROOT output_dir = args.output_dir.expanduser().resolve() + if selected.setup_versions: + rounds = args.rounds if args.rounds is not None else selected.effective_rounds + if args.dry_run: + if args.shuffle_seed is None: + parser.error("--shuffle-seed is required for a setup-matrix dry run") + print( + "\n".join( + plan_lines(selected, shuffle_seed=args.shuffle_seed, first_round=args.first_round, rounds=rounds) + ) + ) + return 0 + if args.shuffle_seed is None or args.qualification is None: + parser.error("--shuffle-seed and --qualification are required unless --dry-run is given") + try: + manifest = run_polars_matrix( + selected, + output_dir=output_dir, + repo_root=repo_root, + shuffle_seed=args.shuffle_seed, + first_round=args.first_round, + rounds=rounds, + qualification=args.qualification.expanduser().resolve(), + ) + except (OSError, RuntimeError, ValueError, subprocess.CalledProcessError) as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + print(f"{selected.display_name} matrix complete: {len(manifest['records'])} runs") + print(f"Manifest: {output_dir / 'matrix-manifest.json'}") + return 0 try: manifest = run_matrix(selected, output_dir=output_dir, repo_root=repo_root, dry_run=args.dry_run) except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as exc: diff --git a/scripts/version_matrix_specs.py b/scripts/version_matrix_specs.py index 437fb80df1..010e55d6c0 100644 --- a/scripts/version_matrix_specs.py +++ b/scripts/version_matrix_specs.py @@ -1,10 +1,29 @@ from __future__ import annotations +import random import re from dataclasses import dataclass PRERELEASE_RE = re.compile(r"\.(?:dev\d+|rc\d+|a\d+|b\d+)$", re.IGNORECASE) +SETUP_ENGINES = {"A": "default", "B": "in-memory", "C": "streaming"} +GIB = 1024**3 + + +@dataclass(frozen=True) +class Cell: + benchmark: str + version: str + setup: str + + @property + def engine(self) -> str: + return SETUP_ENGINES[self.setup] + + @property + def cell_id(self) -> str: + return f"{self.benchmark}/{self.version}/{self.setup}" + def is_prerelease(version: str) -> bool: return bool(PRERELEASE_RE.search(version)) @@ -35,6 +54,20 @@ class MatrixSpec: required_validation: str cleared_database_suffixes: tuple[str, ...] = () driver_option_for_prereleases: bool = False + setup_versions: tuple[tuple[str, tuple[str, ...]], ...] = () + setup_engines: tuple[tuple[str, str], ...] = () + rounds: int | None = None + iterations: int | None = None + thread_count: int | None = None + cell_timeout_s: float | None = None + rss_limit_gib: float | None = None + swap_limit_gib: float | None = None + safety_swap_gib: float | None = None + reference_version: str | None = None + mover_effect: float = 0.10 + qualification_sha256: str | None = None + expected_queries: tuple[tuple[str, int], ...] = () + evidence_validations: tuple[str, ...] = () @property def baseline_version(self) -> str: @@ -69,6 +102,72 @@ def bundle_filename(self, benchmark: str, version: str) -> str: scale = scale_token(self.scales[benchmark]) return f"{benchmark}_sf{scale}_{self.engine}_v{version_token(version)}_median.json" + @property + def effective_rounds(self) -> int: + return self.rounds if self.rounds is not None else self.repetitions + + @property + def setups(self) -> dict[str, tuple[str, ...]]: + return dict(self.setup_versions) + + @property + def setup_engine_map(self) -> dict[str, str]: + return dict(self.setup_engines) + + @property + def query_counts(self) -> dict[str, int]: + return dict(self.expected_queries) + + def engine_for_setup(self, setup: str) -> str: + return dict(self.setup_engines).get(setup, "default") + + def cells(self) -> list[Cell]: + return [ + Cell(benchmark, version, setup) + for benchmark in self.benchmark_ids + for setup, versions in self.setup_versions + for version in versions + ] + + @property + def expected_invocations(self) -> int: + if not self.setup_versions: + return self.expected_benchmark_runs + per_benchmark = sum(len(versions) for _, versions in self.setup_versions) + return self.effective_rounds * per_benchmark * len(self.benchmarks) + + def snapshot(self) -> dict[str, object]: + return { + "platform": self.platform, + "versions": list(self.versions), + "setups": {setup: list(versions) for setup, versions in self.setup_versions}, + "engines": dict(self.setup_engines), + "benchmarks": [{"id": name, "scale": scale} for name, scale in self.benchmarks], + "expected_queries": dict(self.expected_queries), + "rounds": self.effective_rounds, + "iterations": self.iterations, + "thread_count": self.thread_count, + "cell_timeout_s": self.cell_timeout_s, + "rss_limit_gib": self.rss_limit_gib, + "swap_limit_gib": self.swap_limit_gib, + "safety_swap_gib": self.safety_swap_gib, + "reference_version": self.reference_version, + "mover_effect": self.mover_effect, + "qualification_sha256": self.qualification_sha256, + } + + def accepts_without_evidence(self, validation: str | None) -> bool: + return validation == self.required_validation + + def requires_evidence(self, validation: str | None) -> bool: + return validation in self.evidence_validations + + +def round_order(spec: MatrixSpec, shuffle_seed: int, round_index: int) -> list[Cell]: + cells = spec.cells() + random.Random(f"{shuffle_seed}:{round_index}").shuffle(cells) + return cells + DUCKDB = MatrixSpec( engine="duckdb", @@ -83,4 +182,34 @@ def bundle_filename(self, benchmark: str, version: str) -> str: cleared_database_suffixes=(".duckdb", ".wal"), ) -SPECS: dict[str, MatrixSpec] = {spec.engine: spec for spec in (DUCKDB,)} +POLARS = MatrixSpec( + engine="polars", + display_name="Polars", + platform="polars-df", + driver_package="polars", + versions=("1.31.0", "1.35.2", "1.40.1", "1.44.2", "2.0.0"), + benchmarks=(("tpch", 10.0), ("tpcds", 10.0), ("clickbench", 10.0), ("ssb", 10.0)), + repetitions=3, + compression="zstd:3", + required_validation="passed", + setup_versions=( + ("A", ("1.31.0", "1.35.2", "1.40.1", "1.44.2", "2.0.0")), + ("B", ("2.0.0",)), + ("C", ("1.31.0", "1.40.1", "1.44.2", "2.0.0")), + ), + setup_engines=(("A", "default"), ("B", "in-memory"), ("C", "streaming")), + rounds=3, + iterations=1, + thread_count=10, + cell_timeout_s=1800.0, + rss_limit_gib=14.0, + swap_limit_gib=1.0, + safety_swap_gib=6.0, + reference_version="1.44.2", + mover_effect=0.10, + qualification_sha256="da8ea047afdbc880177f6a295dc8d0e9afd1b95a3f044d906e9c6a26b43d2580", + expected_queries=(("tpch", 22), ("tpcds", 103), ("clickbench", 43), ("ssb", 13)), + evidence_validations=("uncertain", "not_run"), +) + +SPECS: dict[str, MatrixSpec] = {spec.engine: spec for spec in (DUCKDB, POLARS)} diff --git a/tests/unit/scripts/test_polars_version_matrix.py b/tests/unit/scripts/test_polars_version_matrix.py index 5593d578ec..022682f930 100644 --- a/tests/unit/scripts/test_polars_version_matrix.py +++ b/tests/unit/scripts/test_polars_version_matrix.py @@ -427,3 +427,76 @@ def test_analyzer_refuses_duplicate_or_missing_rounds(tmp_path: Path) -> None: _analyze(tmp_path, records[:-1]) with pytest.raises(ValueError, match="duplicate record"): _analyze(tmp_path, [*records, records[0]]) + + +def test_polars_entry_points_delegate_to_the_shared_core() -> None: + from scripts import analyze_version_matrix as shared_analyzer, run_version_matrix as shared_runner + from scripts.version_matrix_specs import POLARS as SHARED_POLARS + + assert POLARS is SHARED_POLARS + assert runner.plan_lines is shared_runner.plan_lines + assert runner.cell_command is shared_runner.cell_command + assert runner.parse_freeze is shared_runner.parse_freeze + assert runner.verify_installation is shared_runner.verify_installation + assert runner.check_pins is shared_runner.check_pins + assert runner.classify_outcome is shared_runner.classify_outcome + assert runner.extract_result is shared_runner.extract_result + assert runner.run_matrix is shared_runner.run_polars_matrix + assert analyzer.analyze is shared_analyzer.analyze_polars + assert analyzer.compare is shared_analyzer.compare_cells + assert analyzer.write_outputs is shared_analyzer.write_polars_outputs + assert analyzer.load_qualification is shared_analyzer.load_qualification + + +def test_shared_runner_serves_the_polars_dry_run(capsys: pytest.CaptureFixture[str]) -> None: + from scripts import run_version_matrix as shared_runner + + assert shared_runner.main(["--engine", "polars", "--output-dir", "unused", "--shuffle-seed", "1", "--dry-run"]) == 0 + + assert capsys.readouterr().out.splitlines() == GOLDEN_DRY_RUN.read_text(encoding="utf-8").splitlines() + + +def test_shared_analyzer_serves_the_polars_spec(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from scripts import analyze_version_matrix as shared_analyzer + from scripts.version_matrix_specs import SPECS + + records = [] + for cell in POLARS.cells(): + for index in (1, 2, 3): + record = _record(index, cell.version, cell.setup, 100.0) + record["benchmark"] = cell.benchmark + records.append(record) + rows = [ + {"version": cell.version, "setup": cell.setup, "benchmark": cell.benchmark, "clean": True} + for cell in POLARS.cells() + ] + reference = {benchmark: {"0:1": 4, "0:2": 5} for benchmark in POLARS.benchmark_ids} + qualification_path = tmp_path / "qualification.json" + qualification_path.write_text( + json.dumps({"scale_factor": 10.0, "eligibility": rows, "reference": {"row_counts": reference}}), + encoding="utf-8", + ) + spec = replace(POLARS, qualification_sha256=hashlib.sha256(qualification_path.read_bytes()).hexdigest()) + manifest_path = _write_manifest(tmp_path, spec, records) + monkeypatch.setitem(SPECS, "polars", spec) + output_dir = tmp_path / "out" + + assert ( + shared_analyzer.main( + [ + "--engine", + "polars", + str(manifest_path), + "--qualification", + str(qualification_path), + "--output-dir", + str(output_dir), + ] + ) + == 0 + ) + + result = json.loads((output_dir / "polars-version-matrix-analysis.json").read_text(encoding="utf-8")) + assert result["comparison_count"] == len(result["comparisons"]) + assert all(cell["qualified"] for cell in result["cells"]) + assert (output_dir / "polars-version-matrix-comparisons.csv").is_file()