ci: wait for exact benchmark baselines

.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

.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

.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"
 

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,19 +216,25 @@ 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"])
     if not pr_identities:
         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)

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