From 416cf9349e4e5e2d213de35554e64069e36562de Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 5 Aug 2026 09:46:48 -0400 Subject: [PATCH] ci: wait for exact benchmark baselines Signed-off-by: Connor Tsui --- .github/workflows/pr-bench-runner.yml | 35 +++++-- .github/workflows/pr-bench-sql.yml | 1 + .github/workflows/sql-bench-matrix.yml | 33 ++++-- scripts/compare-benchmark-jsons.py | 62 +++++++++-- scripts/tests/test_benchmark_reporting.py | 34 +++++- scripts/tests/test_wait_for_sql_baseline.py | 108 ++++++++++++++++++++ scripts/wait-for-sql-baseline.sh | 79 ++++++++++++++ 7 files changed, 326 insertions(+), 26 deletions(-) create mode 100644 scripts/tests/test_wait_for_sql_baseline.py create mode 100755 scripts/wait-for-sql-baseline.sh diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index b1b5d8c8799..fda25096528 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -41,7 +41,9 @@ jobs: sccache: s3 - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 with: - ref: ${{ github.event.pull_request.head.sha || github.sha }} + # Benchmark the synthetic merge commit so the measured code includes + # the exact base revision used by this pull_request event. + ref: ${{ github.sha }} - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -78,7 +80,7 @@ jobs: uses: polarsignals/gh-actions-ps-profiling@68ae857e375a826606352016e5b90f01a2a7ff7a # v0.8.1 with: polarsignals_cloud_token: ${{ secrets.POLAR_SIGNALS_API_KEY }} - labels: "branch=${{ github.ref_name }};gh_run_id=${{ github.run_id }};commit_sha=${{ github.event.pull_request.head.sha }};benchmark=${{ inputs.benchmark_id }}" + labels: "branch=${{ github.ref_name }};gh_run_id=${{ github.run_id }};commit_sha=${{ github.sha }};benchmark=${{ inputs.benchmark_id }}" project_uuid: "e5d846e1-b54c-46e7-9174-8bf055a3af56" job_name: "${{ inputs.benchmark_id }}" profiling_frequency: 199 @@ -122,14 +124,33 @@ jobs: - name: Compare results shell: bash + env: + BASELINE_COMMIT: ${{ github.event.pull_request.base.sha }} run: | set -Eeu -o pipefail -x - python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request - gzip -d -c data.json.gz > base.json - - uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ inputs.benchmark_name }}" \ - > comment.md + baseline_args=() + if [[ -n "$BASELINE_COMMIT" ]]; then + baseline_args+=(--baseline-commit "$BASELINE_COMMIT") + fi + + for attempt in {1..40}; do + python3 scripts/s3-download.py \ + s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request + gzip -d -c data.json.gz > base.json + + if uv run --no-project scripts/compare-benchmark-jsons.py \ + base.json results.json "${{ inputs.benchmark_name }}" "${baseline_args[@]}" \ + > comment.md 2> compare.err; then + break + fi + if ! grep -q "No baseline rows found" compare.err || (( attempt == 40 )); then + cat compare.err >&2 + exit 1 + fi + echo "Waiting for benchmark results from base commit $BASELINE_COMMIT" + sleep 60 + done cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR diff --git a/.github/workflows/pr-bench-sql.yml b/.github/workflows/pr-bench-sql.yml index 40789d631a8..67f7e41359d 100644 --- a/.github/workflows/pr-bench-sql.yml +++ b/.github/workflows/pr-bench-sql.yml @@ -31,6 +31,7 @@ on: - "pr-full" permissions: + actions: read # for checking whether the base commit's benchmarks succeeded contents: read pull-requests: write # for commenting on PRs id-token: write # enables AWS-GitHub OIDC diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index 0db85313641..b246aa6c3c9 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -18,16 +18,24 @@ on: jobs: resolve-matrix: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: ${{ inputs.mode == 'pr' && 130 || 10 }} permissions: + actions: read contents: read outputs: benchmark_matrix: ${{ steps.resolve.outputs.benchmark_matrix }} steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 with: - ref: ${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }} + ref: ${{ github.sha }} persist-credentials: false + - name: Wait for SQL baseline + if: inputs.mode == 'pr' && github.event_name == 'pull_request' + shell: bash + env: + BASELINE_COMMIT: ${{ github.event.pull_request.base.sha }} + GH_TOKEN: ${{ github.token }} + run: scripts/wait-for-sql-baseline.sh - name: Install uv uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 with: @@ -59,7 +67,7 @@ jobs: sccache: s3 - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 with: - ref: ${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }} + ref: ${{ github.sha }} - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -120,7 +128,7 @@ jobs: sccache: s3 - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 with: - ref: ${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }} + ref: ${{ github.sha }} - uses: ./.github/actions/setup-rust with: repo-token: ${{ secrets.GITHUB_TOKEN }} @@ -180,7 +188,7 @@ jobs: uses: polarsignals/gh-actions-ps-profiling@68ae857e375a826606352016e5b90f01a2a7ff7a # v0.8.1 with: polarsignals_cloud_token: ${{ secrets.POLAR_SIGNALS_API_KEY }} - labels: "branch=${{ github.ref_name }};gh_run_id=${{ github.run_id }};commit_sha=${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }};benchmark=${{ matrix.id }}" + labels: "branch=${{ github.ref_name }};gh_run_id=${{ github.run_id }};commit_sha=${{ github.sha }};benchmark=${{ matrix.id }}" project_uuid: "e5d846e1-b54c-46e7-9174-8bf055a3af56" job_name: "${{ matrix.id }}" profiling_frequency: 199 @@ -236,20 +244,29 @@ jobs: uv run --no-project scripts/capture-file-sizes.py \ vortex-bench/data \ --benchmark ${{ matrix.subcommand }} \ - --commit ${{ inputs.mode == 'pr' && github.event.pull_request.head.sha || github.sha }} \ + --commit ${{ github.sha }} \ -o sizes.json cat sizes.json >> results.json - name: Compare results if: inputs.mode == 'pr' shell: bash + env: + BASELINE_COMMIT: ${{ github.event.pull_request.base.sha }} run: | set -Eeu -o pipefail -x - python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request + baseline_args=() + if [[ -n "$BASELINE_COMMIT" ]]; then + baseline_args+=(--baseline-commit "$BASELINE_COMMIT") + fi + + python3 scripts/s3-download.py \ + s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request gzip -d -c data.json.gz > base.json - uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.name }}" \ + uv run --no-project scripts/compare-benchmark-jsons.py \ + base.json results.json "${{ matrix.name }}" "${baseline_args[@]}" \ > comment.md cat comment.md >> "$GITHUB_STEP_SUMMARY" diff --git a/scripts/compare-benchmark-jsons.py b/scripts/compare-benchmark-jsons.py index c45b0ba33d6..9d549d1c262 100644 --- a/scripts/compare-benchmark-jsons.py +++ b/scripts/compare-benchmark-jsons.py @@ -11,10 +11,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +import argparse import math import os import re -import sys from dataclasses import dataclass from io import StringIO from typing import Any @@ -216,12 +216,17 @@ def read_jsonl_rows_for_commit(path: str, commit_id: str) -> pd.DataFrame: return pd.DataFrame(rows_by_identity.values()) -def read_latest_baseline_rows(path: str, pr: pd.DataFrame) -> pd.DataFrame: +def read_latest_baseline_rows( + path: str, + pr: pd.DataFrame, + baseline_commit_order: list[str] | None = None, +) -> pd.DataFrame: """Read rows from the latest history commit matching the PR benchmark. A benchmark can be new to the PR workflow and therefore have no baseline yet. Return an empty frame with the PR schema in that case so the report - can show the measurements without comparison. + can show the measurements without comparison. When an exact baseline is + requested, never substitute a different develop commit. """ pr_identities = set(benchmark_identity_rows(pr)["benchmark_identity"]) @@ -229,6 +234,7 @@ def read_latest_baseline_rows(path: str, pr: pd.DataFrame) -> pd.DataFrame: return pd.read_json(path, lines=True) baseline_commit_id = None + matching_commit_ids: set[str] = set() with open(path, encoding="utf-8") as lines: for line in lines: if '"name"' not in line or '"commit_id"' not in line: @@ -238,14 +244,27 @@ def read_latest_baseline_rows(path: str, pr: pd.DataFrame) -> pd.DataFrame: commit_id = record.get("commit_id") if commit_id is not None: baseline_commit_id = commit_id + matching_commit_ids.add(commit_id) + + if baseline_commit_order is not None: + baseline_commit_id = next( + (commit_id for commit_id in baseline_commit_order if commit_id in matching_commit_ids), + None, + ) if baseline_commit_id is None: + if baseline_commit_order is not None: + raise ValueError("No baseline rows found in the requested commit ancestry") return pr.iloc[0:0].copy() return read_jsonl_rows_for_commit(path, baseline_commit_id) -def select_latest_baseline_rows(base: pd.DataFrame, pr: pd.DataFrame) -> pd.DataFrame: +def select_latest_baseline_rows( + base: pd.DataFrame, + pr: pd.DataFrame, + baseline_commit_order: list[str] | None = None, +) -> pd.DataFrame: """Select rows from the latest baseline commit containing this benchmark. The persisted benchmark history is append-only. A row only appears after @@ -258,6 +277,10 @@ def select_latest_baseline_rows(base: pd.DataFrame, pr: pd.DataFrame) -> pd.Data commit_ids = base["commit_id"].dropna().unique() if len(commit_ids) <= 1: + if baseline_commit_order is not None and any( + commit_id not in baseline_commit_order for commit_id in commit_ids + ): + raise ValueError("No baseline rows found in the requested commit ancestry") return base pr_identities = set(benchmark_identity_rows(pr)["benchmark_identity"]) @@ -270,7 +293,16 @@ def select_latest_baseline_rows(base: pd.DataFrame, pr: pd.DataFrame) -> pd.Data if matches.empty: return base.iloc[0:0].copy() - baseline_commit_id = matches["commit_id"].iloc[-1] + if baseline_commit_order is None: + baseline_commit_id = matches["commit_id"].iloc[-1] + else: + matching_commit_ids = set(matches["commit_id"]) + baseline_commit_id = next( + (commit_id for commit_id in baseline_commit_order if commit_id in matching_commit_ids), + None, + ) + if baseline_commit_id is None: + raise ValueError("No baseline rows found in the requested commit ancestry") return base[base["commit_id"] == baseline_commit_id].copy() @@ -957,11 +989,21 @@ def group_sort_key(group_key: tuple[str, str, str]) -> tuple[int, int, int, str, def main() -> None: """Render the benchmark comparison markdown used in CI PR comments.""" - benchmark_name = sys.argv[3] if len(sys.argv) > 3 else "" + parser = argparse.ArgumentParser() + parser.add_argument("base_path") + parser.add_argument("pr_path") + parser.add_argument("benchmark_name", nargs="?", default="") + parser.add_argument( + "--baseline-commit", + help="Exact develop commit whose benchmark results must be used", + ) + args = parser.parse_args() + + baseline_commit_order = [args.baseline_commit] if args.baseline_commit else None - pr = pd.read_json(sys.argv[2], lines=True) - title = format_title(benchmark_name, pr) - base = read_latest_baseline_rows(sys.argv[1], pr) + pr = pd.read_json(args.pr_path, lines=True) + title = format_title(args.benchmark_name, pr) + base = read_latest_baseline_rows(args.base_path, pr, baseline_commit_order) base_commit_ids = set(base["commit_id"].unique()) pr_commit_id = set(pr["commit_id"].unique()) @@ -981,7 +1023,7 @@ def main() -> None: df3["unit"] = df3["unit"].fillna("unit") df3["ratio"] = df3["value_pr"] / df3["value_base"] - is_s3_benchmark = "s3" in benchmark_name.lower() + is_s3_benchmark = "s3" in args.benchmark_name.lower() threshold_pct = 30 if is_s3_benchmark else 10 improvement_threshold = 1.0 - (threshold_pct / 100.0) regression_threshold = 1.0 + (threshold_pct / 100.0) diff --git a/scripts/tests/test_benchmark_reporting.py b/scripts/tests/test_benchmark_reporting.py index 8f187c3b747..54ca58f924d 100644 --- a/scripts/tests/test_benchmark_reporting.py +++ b/scripts/tests/test_benchmark_reporting.py @@ -8,6 +8,7 @@ from pathlib import Path import pandas as pd +import pytest REPO_ROOT = Path(__file__).resolve().parents[2] COMPARE_SCRIPT = REPO_ROOT / "scripts" / "compare-benchmark-jsons.py" @@ -202,6 +203,37 @@ def test_read_latest_baseline_rows_streams_latest_matching_benchmark_commit(tmp_ assert len(selected) == 2 +def test_read_latest_baseline_rows_uses_requested_base_commit(tmp_path: Path) -> None: + compare = load_compare_module() + history_path = tmp_path / "history.jsonl" + history_rows = [ + stored_timing_row("base-ancestor", "tpch_q01/datafusion:parquet", 100), + stored_timing_row("future-develop", "tpch_q01/datafusion:parquet", 90), + ] + history_path.write_text( + "".join(f"{json.dumps(row)}\n" for row in history_rows), + encoding="utf-8", + ) + pr = pd.DataFrame([stored_timing_row("pr-sha", "tpch_q01/datafusion:parquet", 105)]) + + selected = compare.read_latest_baseline_rows(history_path, pr, ["base-ancestor"]) + + assert set(selected["commit_id"]) == {"base-ancestor"} + + +def test_read_latest_baseline_rows_does_not_fall_back_from_requested_commit(tmp_path: Path) -> None: + compare = load_compare_module() + history_path = tmp_path / "history.jsonl" + history_path.write_text( + f"{json.dumps(stored_timing_row('other-commit', 'tpch_q01/datafusion:parquet', 100))}\n", + encoding="utf-8", + ) + pr = pd.DataFrame([stored_timing_row("pr-sha", "tpch_q01/datafusion:parquet", 105)]) + + with pytest.raises(ValueError, match="No baseline rows found"): + compare.read_latest_baseline_rows(history_path, pr, ["requested-base"]) + + def test_read_latest_baseline_rows_uses_last_result_from_rerun(tmp_path: Path) -> None: compare = load_compare_module() history_path = tmp_path / "history.jsonl" @@ -217,7 +249,7 @@ def test_read_latest_baseline_rows_uses_last_result_from_rerun(tmp_path: Path) - ) pr = pd.DataFrame([stored_timing_row("pr", "tpch_q01/datafusion:parquet", 105)]) - selected = compare.read_latest_baseline_rows(history_path, pr) + selected = compare.read_latest_baseline_rows(history_path, pr, ["base"]) assert len(selected) == 1 assert selected.iloc[0]["value"] == 110 diff --git a/scripts/tests/test_wait_for_sql_baseline.py b/scripts/tests/test_wait_for_sql_baseline.py new file mode 100644 index 00000000000..646790cc94f --- /dev/null +++ b/scripts/tests/test_wait_for_sql_baseline.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import json +import os +import subprocess +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +WAIT_SCRIPT = REPO_ROOT / "scripts" / "wait-for-sql-baseline.sh" + + +def run_wait_script( + tmp_path: Path, + run_status: str, + jobs: list[dict[str, str | None]], +) -> subprocess.CompletedProcess[str]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + + gh = bin_dir / "gh" + gh.write_text( + """#!/usr/bin/env bash +if [[ "$*" == *"/jobs"* ]]; then + printf '%s\n' "$MOCK_JOBS" +else + printf '{"workflow_runs":[{"id":42,"path":".github/workflows/develop-bench.yml","status":"%s"}]}\n' \\ + "$MOCK_RUN_STATUS" +fi +""", + encoding="utf-8", + ) + gh.chmod(0o755) + + sleep = bin_dir / "sleep" + sleep.write_text("#!/usr/bin/env bash\nexit 0\n", encoding="utf-8") + sleep.chmod(0o755) + + env = os.environ.copy() + env.update( + { + "BASELINE_COMMIT": "base-commit", + "GITHUB_REPOSITORY": "vortex-data/vortex", + "MOCK_JOBS": json.dumps({"jobs": jobs}), + "MOCK_RUN_STATUS": run_status, + "PATH": f"{bin_dir}:{env['PATH']}", + } + ) + return subprocess.run( + ["bash", str(WAIT_SCRIPT)], + check=False, + capture_output=True, + text=True, + env=env, + ) + + +def sql_job(name: str, status: str, conclusion: str | None) -> dict[str, str | None]: + return { + "name": f"sql / bench ({name})", + "status": status, + "conclusion": conclusion, + } + + +def test_succeeds_when_sql_jobs_finish_before_workflow(tmp_path: Path) -> None: + result = run_wait_script( + tmp_path, + "in_progress", + [ + sql_job("tpch", "completed", "success"), + sql_job("clickbench", "completed", "success"), + ], + ) + + assert result.returncode == 0 + assert "SQL baselines are ready for base-commit" in result.stdout + + +def test_waits_while_any_sql_job_is_incomplete(tmp_path: Path) -> None: + result = run_wait_script( + tmp_path, + "in_progress", + [ + sql_job("tpch", "completed", "success"), + sql_job("clickbench", "in_progress", None), + ], + ) + + assert result.returncode == 1 + assert "Timed out waiting for SQL baselines from base-commit" in result.stderr + assert "SQL benchmarks failed" not in result.stderr + + +def test_fails_immediately_when_sql_job_fails(tmp_path: Path) -> None: + result = run_wait_script( + tmp_path, + "in_progress", + [ + sql_job("tpch", "completed", "failure"), + sql_job("clickbench", "in_progress", None), + ], + ) + + assert result.returncode == 1 + assert "SQL benchmarks failed for base commit base-commit" in result.stderr + assert "sql / bench (tpch): failure" in result.stderr + assert "clickbench" not in result.stderr diff --git a/scripts/wait-for-sql-baseline.sh b/scripts/wait-for-sql-baseline.sh new file mode 100755 index 00000000000..1194713ecc0 --- /dev/null +++ b/scripts/wait-for-sql-baseline.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +# Wait for successful SQL benchmark results for the pull request base commit. +set -Eeuo pipefail + +max_attempts=120 +for ((attempt = 1; attempt <= max_attempts; attempt++)); do + runs="$( + gh api --method GET "repos/${GITHUB_REPOSITORY}/actions/runs" \ + -f head_sha="$BASELINE_COMMIT" \ + -f event=push \ + -f per_page=100 + )" + run_id="$( + jq -r ' + [.workflow_runs[] + | select( + .path == ".github/workflows/bench.yml" + or .path == ".github/workflows/develop-bench.yml" + )] + | max_by(.id) + | .id // empty + ' <<< "$runs" + )" + + if [[ -n "$run_id" ]]; then + run_status="$(jq -r --argjson run_id "$run_id" ' + .workflow_runs[] + | select(.id == $run_id) + | .status + ' <<< "$runs")" + + jobs="$( + gh api --method GET \ + "repos/${GITHUB_REPOSITORY}/actions/runs/${run_id}/jobs" \ + -f filter=latest \ + -f per_page=100 + )" + sql_jobs="$(jq '[.jobs[] | select(.name | startswith("sql / bench ("))]' <<< "$jobs")" + sql_job_count="$(jq length <<< "$sql_jobs")" + failed_jobs="$( + jq '[.[] | select(.status == "completed" and .conclusion != "success")] | length' \ + <<< "$sql_jobs" + )" + incomplete_jobs="$(jq '[.[] | select(.status != "completed")] | length' <<< "$sql_jobs")" + + if (( failed_jobs > 0 )); then + echo "SQL benchmarks failed for base commit $BASELINE_COMMIT:" >&2 + jq -r ' + .[] + | select(.status == "completed" and .conclusion != "success") + | " \(.name): \(.conclusion)" + ' \ + <<< "$sql_jobs" >&2 + exit 1 + fi + + if (( sql_job_count > 0 && incomplete_jobs == 0 )); then + echo "SQL baselines are ready for $BASELINE_COMMIT" + exit 0 + fi + + if [[ "$run_status" == "completed" ]]; then + echo "Develop workflow completed without SQL baselines for $BASELINE_COMMIT" >&2 + exit 1 + fi + fi + + if (( attempt == max_attempts )); then + echo "Timed out waiting for SQL baselines from $BASELINE_COMMIT" >&2 + exit 1 + fi + + echo "Waiting for SQL benchmarks from base commit $BASELINE_COMMIT" + sleep 60 +done