From a2ab0f69d48012a6ff545065623dbfb2ffb0c899 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Fri, 17 Jul 2026 23:33:13 +0530
Subject: [PATCH 1/3] Add automated benchmark regression check in CI (#144)

Run the scaling benchmark on the PR and its base branch back-to-back on
the same runner, swapping only the library under test (the src/ layout
keeps the PR's benchmark harness fixed), and post a sticky pull-request
comment with the relative runtime change per grid size. Compares relative
(%) change rather than absolute seconds so per-runner noise largely
cancels. Informational and non-blocking for now, with a low starting
threshold to be calibrated against the runner's noise floor.

- tests/benchmarks/compare.py: parse two --output-json files (schema from
  #140), compute per-grid-size % deltas, render the markdown table, flag
  above a configurable threshold
- tests/benchmarks/test_compare.py: unit tests for the comparison logic
- .github/workflows/benchmark-regression.yml: same-runner A/B workflow,
  sticky comment via actions/github-script (job summary fallback for forks)
---
 .github/workflows/benchmark-regression.yml | 112 +++++++++
 CHANGELOG.md                               |   7 +
 tests/benchmarks/compare.py                | 279 +++++++++++++++++++++
 tests/benchmarks/test_compare.py           | 163 ++++++++++++
 4 files changed, 561 insertions(+)
 create mode 100644 .github/workflows/benchmark-regression.yml
 create mode 100644 tests/benchmarks/compare.py
 create mode 100644 tests/benchmarks/test_compare.py

diff --git a/.github/workflows/benchmark-regression.yml b/.github/workflows/benchmark-regression.yml
new file mode 100644
index 0000000..ad32b79
--- /dev/null
+++ b/.github/workflows/benchmark-regression.yml
@@ -0,0 +1,112 @@
+# Automated benchmark regression check (PR vs base branch).
+#
+# Design (see https://github.com/mllam/weather-model-graphs/issues/144):
+#   * Run the scaling benchmark twice, back-to-back, on the SAME runner: once
+#     against the PR's install of the library and once against the base branch's
+#     install. Only the library under test is swapped - the benchmark harness
+#     itself always comes from the PR checkout (the package uses a src/ layout,
+#     so the working tree never shadows the installed library).
+#   * Compare RELATIVE (%) change, not absolute seconds, so per-runner noise
+#     largely cancels.
+#   * Post the result as a single, updating ("sticky") PR comment. The check is
+#     informational and NON-BLOCKING for now; the threshold starts low and can
+#     be raised once the runner's real noise floor is known.
+
+name: benchmark regression
+
+on:
+  pull_request:
+    paths:
+      - "src/weather_model_graphs/create/**"
+      - "tests/benchmarks/**"
+      - ".github/workflows/benchmark-regression.yml"
+  workflow_dispatch:
+
+# Only the latest run per PR needs to finish; cancel superseded runs.
+concurrency:
+  group: benchmark-regression-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    env:
+      # Kept small so the whole job stays within a few minutes. Tune here.
+      BENCH_ARGS: "--min-N 50 --max-N 200 --num-steps 4 --archetype keisler"
+      # Start low; raise once we've seen the runner's noise floor across a few
+      # real runs (per the discussion on #144).
+      THRESHOLD_PCT: "0.1"
+      BASE_REF: ${{ github.event.pull_request.base.ref || 'main' }}
+    steps:
+      - name: Checkout PR
+        uses: actions/checkout@v5
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.12"
+
+      - name: Run benchmark on PR and base, then compare
+        run: |
+          set -euo pipefail
+          uv venv
+          source .venv/bin/activate
+
+          echo "::group::Install PR build and benchmark it"
+          uv pip install ".[visualisation]"
+          python -m tests.benchmarks.graph_creation_scaling $BENCH_ARGS \
+            --output-json pr.json --output-plot-runtime pr_plot.png
+          echo "::endgroup::"
+
+          echo "::group::Swap in the '$BASE_REF' library and benchmark it"
+          # Reinstall ONLY the library from the base branch; the harness on disk
+          # (the PR's) is untouched, so we measure with one fixed ruler.
+          uv pip install --reinstall-package weather-model-graphs \
+            "weather-model-graphs @ git+https://github.com/mllam/weather-model-graphs.git@${BASE_REF}"
+          python -m tests.benchmarks.graph_creation_scaling $BENCH_ARGS \
+            --output-json base.json --output-plot-runtime base_plot.png
+          echo "::endgroup::"
+
+          python -m tests.benchmarks.compare base.json pr.json \
+            --threshold-pct "$THRESHOLD_PCT" \
+            --baseline-label "$BASE_REF" --contender-label "PR" \
+            --output comment.md
+
+      - name: Publish result to the job summary
+        if: always()
+        run: cat comment.md >> "$GITHUB_STEP_SUMMARY" || true
+
+      - name: Post or update sticky PR comment
+        # Only same-repo PRs get a writable token; fork PRs fall back to the job
+        # summary above. Never fail the job over the comment (informational).
+        if: github.event_name == 'pull_request'
+        continue-on-error: true
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const fs = require('fs');
+            const marker = '<!-- benchmark-regression-check -->';
+            const body = fs.readFileSync('comment.md', 'utf8');
+            const { owner, repo } = context.repo;
+            const issue_number = context.issue.number;
+            const comments = await github.paginate(
+              github.rest.issues.listComments,
+              { owner, repo, issue_number, per_page: 100 }
+            );
+            const existing = comments.find(
+              (c) => c.body && c.body.includes(marker)
+            );
+            if (existing) {
+              await github.rest.issues.updateComment({
+                owner, repo, comment_id: existing.id, body,
+              });
+            } else {
+              await github.rest.issues.createComment({
+                owner, repo, issue_number, body,
+              });
+            }
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3794f22..59b063d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -32,6 +32,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add support for writing benchmarking results to json,
   [\#140](https://github.com/mllam/weather-model-graphs/pull/140),
   @yuvraajnarula & @leifdenby
+- Add an automated benchmark regression check in CI: the scaling benchmark is
+  run on the PR and its base branch back-to-back on the same runner (swapping
+  only the library under test), and a sticky pull-request comment reports the
+  relative runtime change per grid size. Adds `tests/benchmarks/compare.py` and
+  a GitHub Actions workflow; informational and non-blocking for now, with a low
+  starting threshold to be calibrated against the runner's noise floor.
+  [\#144](https://github.com/mllam/weather-model-graphs/issues/144), @prajwal-tech07
 
 ### Deprecated
 
diff --git a/tests/benchmarks/compare.py b/tests/benchmarks/compare.py
new file mode 100644
index 0000000..dc03684
--- /dev/null
+++ b/tests/benchmarks/compare.py
@@ -0,0 +1,279 @@
+"""Compare two scaling-benchmark JSON outputs and render a regression report.
+
+This consumes the JSON produced by ``graph_creation_scaling.py --output-json``
+(added in #140), i.e. a list of records of the form::
+
+    [{"grid_points": int, "runtime_s": float, "peak_memory_mb": float | null}, ...]
+
+Given a *baseline* run (typically ``main``) and a *contender* run (the PR), it
+matches records by ``grid_points``, computes the relative change in runtime, and
+renders a Markdown table suitable for posting as a sticky pull-request comment.
+
+The design follows the discussion in
+https://github.com/mllam/weather-model-graphs/issues/144: we compare relative
+(%) change rather than absolute seconds, because the two runs are executed
+back-to-back on the *same* CI runner and only the library under test differs, so
+the per-runner noise largely cancels.
+
+Usage::
+
+    python -m tests.benchmarks.compare main.json pr.json
+    python -m tests.benchmarks.compare main.json pr.json \\
+        --threshold-pct 0.1 --output comment.md --fail-on-regression
+
+The script only depends on the standard library so it can run in a minimal CI
+step without installing the package.
+"""
+
+import argparse
+import json
+import sys
+from typing import Dict, List, Optional
+
+# Hidden marker used by the CI workflow to find-and-update a single sticky
+# comment instead of posting a new comment on every run.
+STICKY_MARKER = "<!-- benchmark-regression-check -->"
+
+
+def load_results(path: str) -> Dict[int, dict]:
+    """Load a benchmark JSON file and index the records by ``grid_points``.
+
+    Raises ``ValueError`` with an actionable message if the file is empty or
+    does not match the expected schema, so CI failures are easy to diagnose.
+    """
+    with open(path) as f:
+        data = json.load(f)
+
+    if not isinstance(data, list) or not data:
+        raise ValueError(
+            f"{path!r} does not contain a non-empty list of benchmark records"
+        )
+
+    indexed: Dict[int, dict] = {}
+    for record in data:
+        if "grid_points" not in record or "runtime_s" not in record:
+            raise ValueError(
+                f"{path!r} contains a record missing required keys "
+                f"'grid_points'/'runtime_s': {record!r}"
+            )
+        indexed[int(record["grid_points"])] = record
+    return indexed
+
+
+def _pct_change(baseline: float, contender: float) -> Optional[float]:
+    """Return the percentage change from ``baseline`` to ``contender``.
+
+    Returns ``None`` when the baseline is zero (or negative), since a relative
+    change is undefined there and we would rather skip the row than divide by
+    zero.
+    """
+    if baseline <= 0:
+        return None
+    return (contender - baseline) / baseline * 100.0
+
+
+class Row:
+    """A single grid-size comparison line in the report."""
+
+    def __init__(
+        self,
+        grid_points: int,
+        baseline_s: float,
+        contender_s: float,
+        delta_pct: Optional[float],
+        is_regression: bool,
+    ):
+        self.grid_points = grid_points
+        self.baseline_s = baseline_s
+        self.contender_s = contender_s
+        self.delta_pct = delta_pct
+        self.is_regression = is_regression
+
+
+def compare(
+    baseline: Dict[int, dict],
+    contender: Dict[int, dict],
+    threshold_pct: float,
+) -> List[Row]:
+    """Build the per-grid-size comparison rows for the runtime metric.
+
+    Only ``grid_points`` present in *both* runs are compared; unmatched sizes
+    are skipped (they show up in the PR diff of the benchmark itself, so there is
+    nothing to compare against). Rows are returned sorted by ``grid_points``.
+    """
+    common = sorted(set(baseline) & set(contender))
+    rows: List[Row] = []
+    for gp in common:
+        b = float(baseline[gp]["runtime_s"])
+        c = float(contender[gp]["runtime_s"])
+        delta = _pct_change(b, c)
+        is_regression = delta is not None and delta > threshold_pct
+        rows.append(Row(gp, b, c, delta, is_regression))
+    return rows
+
+
+def _fmt_seconds(value: float) -> str:
+    """Format a runtime in seconds with millisecond-level readability."""
+    if value < 1.0:
+        return f"{value * 1000:.0f}ms"
+    return f"{value:.3f}s"
+
+
+def _fmt_delta(delta: Optional[float], is_regression: bool) -> str:
+    if delta is None:
+        return "n/a"
+    icon = "⚠️" if is_regression else "✅"
+    return f"{delta:+.1f}% {icon}"
+
+
+def render_markdown(
+    rows: List[Row],
+    threshold_pct: float,
+    baseline_label: str,
+    contender_label: str,
+    unmatched: Optional[List[int]] = None,
+) -> str:
+    """Render the comparison as a Markdown block, prefixed with the sticky marker."""
+    lines = [STICKY_MARKER, "## ⏱️ Graph-creation benchmark: regression check", ""]
+
+    if not rows:
+        lines.append(
+            "No overlapping grid sizes to compare between "
+            f"`{baseline_label}` and `{contender_label}`."
+        )
+        return "\n".join(lines) + "\n"
+
+    regressions = [r for r in rows if r.is_regression]
+    if regressions:
+        lines.append(
+            f"⚠️ **{len(regressions)} of {len(rows)}** grid sizes exceed the "
+            f"**+{threshold_pct:g}%** runtime threshold."
+        )
+    else:
+        lines.append(
+            f"✅ No runtime regression above **+{threshold_pct:g}%** "
+            f"across {len(rows)} grid sizes."
+        )
+    lines.append("")
+
+    lines.append(f"| grid points | {baseline_label} | {contender_label} | Δ runtime |")
+    lines.append("|---:|---:|---:|---:|")
+    for r in rows:
+        lines.append(
+            f"| {r.grid_points:,} | {_fmt_seconds(r.baseline_s)} | "
+            f"{_fmt_seconds(r.contender_s)} | {_fmt_delta(r.delta_pct, r.is_regression)} |"
+        )
+
+    if unmatched:
+        pretty = ", ".join(f"{gp:,}" for gp in unmatched)
+        lines.append("")
+        lines.append(
+            f"_Note: {len(unmatched)} grid size(s) not present in both runs were "
+            f"skipped: {pretty}._"
+        )
+
+    lines.append("")
+    lines.append(
+        "_Runs execute back-to-back on the same runner; only the library under "
+        "test differs, so relative change is compared rather than absolute time._"
+    )
+    return "\n".join(lines) + "\n"
+
+
+def build_report(
+    baseline_path: str,
+    contender_path: str,
+    threshold_pct: float,
+    baseline_label: str,
+    contender_label: str,
+):
+    """Load both files, compute rows, and return ``(markdown, rows)``."""
+    baseline = load_results(baseline_path)
+    contender = load_results(contender_path)
+    rows = compare(baseline, contender, threshold_pct)
+    unmatched = sorted(set(baseline) ^ set(contender))
+    markdown = render_markdown(
+        rows, threshold_pct, baseline_label, contender_label, unmatched
+    )
+    return markdown, rows
+
+
+def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compare two scaling-benchmark JSON outputs (baseline vs PR)."
+    )
+    parser.add_argument("baseline", help="Baseline JSON (e.g. main.json).")
+    parser.add_argument("contender", help="Contender JSON (e.g. pr.json).")
+    parser.add_argument(
+        "--threshold-pct",
+        type=float,
+        default=0.1,
+        help="Flag a grid size when its runtime increases by more than this "
+        "percentage. Start low and raise it once the runner's noise floor is "
+        "known (default: 0.1).",
+    )
+    parser.add_argument(
+        "--baseline-label", default="main", help="Column label for the baseline."
+    )
+    parser.add_argument(
+        "--contender-label", default="PR", help="Column label for the contender."
+    )
+    parser.add_argument(
+        "--output",
+        help="Also write the Markdown report to this file (for the CI comment).",
+    )
+    parser.add_argument(
+        "--fail-on-regression",
+        action="store_true",
+        help="Exit non-zero if any grid size regresses (off by default so the "
+        "check is informational).",
+    )
+    return parser.parse_args(argv)
+
+
+def _print_utf8(text: str) -> None:
+    """Print ``text`` as UTF-8 regardless of the console's default encoding.
+
+    The report contains emoji (✅/⚠️); on a Windows console (cp1252) a plain
+    ``print`` raises ``UnicodeEncodeError``. CI runners are UTF-8, but we keep
+    the tool robust for local use.
+    """
+    stream = sys.stdout
+    reconfigure = getattr(stream, "reconfigure", None)
+    if reconfigure is not None:
+        try:
+            reconfigure(encoding="utf-8")
+        except (ValueError, OSError):
+            pass
+    try:
+        print(text)
+    except UnicodeEncodeError:
+        buffer = getattr(stream, "buffer", None)
+        if buffer is not None:
+            buffer.write(text.encode("utf-8") + b"\n")
+        else:  # pragma: no cover - extremely unusual stdout replacement
+            print(text.encode("utf-8", "backslashreplace").decode("ascii"))
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    args = parse_args(argv)
+    markdown, rows = build_report(
+        args.baseline,
+        args.contender,
+        args.threshold_pct,
+        args.baseline_label,
+        args.contender_label,
+    )
+
+    _print_utf8(markdown)
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as f:
+            f.write(markdown)
+
+    if args.fail_on_regression and any(r.is_regression for r in rows):
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tests/benchmarks/test_compare.py b/tests/benchmarks/test_compare.py
new file mode 100644
index 0000000..b7852cd
--- /dev/null
+++ b/tests/benchmarks/test_compare.py
@@ -0,0 +1,163 @@
+"""Unit tests for the benchmark comparison script (``tests/benchmarks/compare.py``).
+
+These are pure-stdlib tests (no ``weather_model_graphs`` import needed) so they
+run in the ordinary pytest job and give us confidence in the regression logic
+independently of an actual benchmark run.
+"""
+
+import json
+
+import pytest
+
+from tests.benchmarks import compare
+
+
+def _write(tmp_path, name, records):
+    path = tmp_path / name
+    path.write_text(json.dumps(records))
+    return str(path)
+
+
+def _rec(grid_points, runtime_s, peak_memory_mb=None):
+    return {
+        "grid_points": grid_points,
+        "runtime_s": runtime_s,
+        "peak_memory_mb": peak_memory_mb,
+    }
+
+
+def test_pct_change_basic():
+    assert compare._pct_change(1.0, 1.5) == pytest.approx(50.0)
+    assert compare._pct_change(2.0, 1.0) == pytest.approx(-50.0)
+
+
+def test_pct_change_zero_baseline_is_none():
+    assert compare._pct_change(0.0, 1.0) is None
+    assert compare._pct_change(-1.0, 1.0) is None
+
+
+def test_compare_flags_regression_above_threshold():
+    baseline = {1024: _rec(1024, 1.00)}
+    contender = {1024: _rec(1024, 1.30)}  # +30%
+    rows = compare.compare(baseline, contender, threshold_pct=25.0)
+    assert len(rows) == 1
+    assert rows[0].delta_pct == pytest.approx(30.0)
+    assert rows[0].is_regression is True
+
+
+def test_compare_no_flag_below_threshold():
+    baseline = {1024: _rec(1024, 1.00)}
+    contender = {1024: _rec(1024, 1.02)}  # +2%
+    rows = compare.compare(baseline, contender, threshold_pct=25.0)
+    assert rows[0].is_regression is False
+
+
+def test_threshold_boundary_is_strict_greater_than():
+    # Exactly at threshold must NOT flag (we only flag strictly above it).
+    # Use an exact 0% delta against a 0% threshold to avoid float-rounding fuzz.
+    baseline = {1024: _rec(1024, 1.00)}
+    contender = {1024: _rec(1024, 1.00)}  # +0.0%
+    rows = compare.compare(baseline, contender, threshold_pct=0.0)
+    assert rows[0].delta_pct == pytest.approx(0.0)
+    assert rows[0].is_regression is False  # 0.0 is not > 0.0
+
+    # And a hair above the threshold must flag.
+    rows_above = compare.compare(
+        {1024: _rec(1024, 1.00)}, {1024: _rec(1024, 1.05)}, threshold_pct=1.0
+    )
+    assert rows_above[0].is_regression is True
+
+
+def test_improvement_is_never_a_regression():
+    baseline = {1024: _rec(1024, 2.00)}
+    contender = {1024: _rec(1024, 1.00)}  # -50%
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    assert rows[0].delta_pct == pytest.approx(-50.0)
+    assert rows[0].is_regression is False
+
+
+def test_only_common_grid_points_compared_and_sorted():
+    baseline = {4096: _rec(4096, 2.0), 1024: _rec(1024, 1.0), 256: _rec(256, 0.5)}
+    contender = {1024: _rec(1024, 1.0), 4096: _rec(4096, 2.0), 9001: _rec(9001, 9.0)}
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    assert [r.grid_points for r in rows] == [1024, 4096]  # sorted, intersection only
+
+
+def test_zero_baseline_row_is_not_a_regression():
+    baseline = {1024: _rec(1024, 0.0)}
+    contender = {1024: _rec(1024, 1.0)}
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    assert rows[0].delta_pct is None
+    assert rows[0].is_regression is False
+
+
+def test_load_results_rejects_empty(tmp_path):
+    path = _write(tmp_path, "empty.json", [])
+    with pytest.raises(ValueError, match="non-empty list"):
+        compare.load_results(path)
+
+
+def test_load_results_rejects_missing_keys(tmp_path):
+    path = _write(tmp_path, "bad.json", [{"grid_points": 10}])
+    with pytest.raises(ValueError, match="missing required keys"):
+        compare.load_results(path)
+
+
+def test_load_results_indexes_by_grid_points(tmp_path):
+    path = _write(tmp_path, "ok.json", [_rec(1024, 1.0), _rec(4096, 2.0)])
+    indexed = compare.load_results(path)
+    assert set(indexed) == {1024, 4096}
+    assert indexed[4096]["runtime_s"] == 2.0
+
+
+def test_render_markdown_contains_marker_and_table():
+    rows = compare.compare(
+        {1024: _rec(1024, 1.0)}, {1024: _rec(1024, 1.3)}, threshold_pct=0.1
+    )
+    md = compare.render_markdown(rows, 0.1, "main", "PR")
+    assert compare.STICKY_MARKER in md
+    assert "| grid points | main | PR | Δ runtime |" in md
+    assert "+30.0% ⚠️" in md
+    assert "1,024" in md  # thousands separator
+
+
+def test_render_markdown_clean_run_reports_pass():
+    rows = compare.compare(
+        {1024: _rec(1024, 1.0)}, {1024: _rec(1024, 1.0)}, threshold_pct=0.1
+    )
+    md = compare.render_markdown(rows, 0.1, "main", "PR")
+    assert "No runtime regression" in md
+    assert "⚠️" not in md
+
+
+def test_render_markdown_no_overlap_message():
+    md = compare.render_markdown([], 0.1, "main", "PR", unmatched=[1024])
+    assert "No overlapping grid sizes" in md
+
+
+def test_build_report_end_to_end(tmp_path):
+    base = _write(tmp_path, "main.json", [_rec(1024, 1.0), _rec(4096, 2.0)])
+    cont = _write(tmp_path, "pr.json", [_rec(1024, 1.0), _rec(4096, 3.0)])  # +50%
+    md, rows = compare.build_report(base, cont, 0.1, "main", "PR")
+    assert len(rows) == 2
+    assert any(r.is_regression for r in rows)
+    assert "+50.0% ⚠️" in md
+
+
+def test_main_fail_on_regression_exit_code(tmp_path, capsys):
+    base = _write(tmp_path, "main.json", [_rec(1024, 1.0)])
+    cont = _write(tmp_path, "pr.json", [_rec(1024, 2.0)])
+    rc = compare.main([base, cont, "--threshold-pct", "0.1", "--fail-on-regression"])
+    assert rc == 1
+    # Without the flag it must stay informational (exit 0).
+    rc2 = compare.main([base, cont, "--threshold-pct", "0.1"])
+    assert rc2 == 0
+
+
+def test_main_writes_output_file(tmp_path):
+    base = _write(tmp_path, "main.json", [_rec(1024, 1.0)])
+    cont = _write(tmp_path, "pr.json", [_rec(1024, 1.0)])
+    out = tmp_path / "comment.md"
+    rc = compare.main([base, cont, "--output", str(out)])
+    assert rc == 0
+    assert compare.STICKY_MARKER in out.read_text(encoding="utf-8")

From 90f918d4a417b4987e14298d3a15b4928badd93d Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Thu, 13 Aug 2026 23:33:45 +0530
Subject: [PATCH 2/3] Use marocchino/sticky-pull-request-comment for the
 benchmark comment

Replace the hand-rolled actions/github-script step (list comments,
find-by-marker, create-or-update) with the purpose-built action, per
review feedback on #147. Behaviour is unchanged: same comment.md input,
same fork fallback to the job summary, same non-blocking failure mode.
---
 .github/workflows/benchmark-regression.yml | 26 +++-------------------
 1 file changed, 3 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/benchmark-regression.yml b/.github/workflows/benchmark-regression.yml
index ad32b79..a26d17a 100644
--- a/.github/workflows/benchmark-regression.yml
+++ b/.github/workflows/benchmark-regression.yml
@@ -86,27 +86,7 @@ jobs:
         # summary above. Never fail the job over the comment (informational).
         if: github.event_name == 'pull_request'
         continue-on-error: true
-        uses: actions/github-script@v7
+        uses: marocchino/sticky-pull-request-comment@v3
         with:
-          script: |
-            const fs = require('fs');
-            const marker = '<!-- benchmark-regression-check -->';
-            const body = fs.readFileSync('comment.md', 'utf8');
-            const { owner, repo } = context.repo;
-            const issue_number = context.issue.number;
-            const comments = await github.paginate(
-              github.rest.issues.listComments,
-              { owner, repo, issue_number, per_page: 100 }
-            );
-            const existing = comments.find(
-              (c) => c.body && c.body.includes(marker)
-            );
-            if (existing) {
-              await github.rest.issues.updateComment({
-                owner, repo, comment_id: existing.id, body,
-              });
-            } else {
-              await github.rest.issues.createComment({
-                owner, repo, issue_number, body,
-              });
-            }
+          header: benchmark-regression
+          path: comment.md

From d14c2f784d191e84270e64e7aecbcf0961505ea7 Mon Sep 17 00:00:00 2001
From: prajwal <prajwalhawaldar2@gmail.com>
Date: Fri, 14 Aug 2026 00:05:15 +0530
Subject: [PATCH 3/3] Include peak-memory delta in the benchmark regression
 check (Phase 1)

Per discussion on #144: memory usage moves into Phase 1 (Phase 2 is now
repetitions/median only). compare.py now reports a peak-memory delta
column alongside runtime whenever both runs recorded --track-memory
data, using the same threshold/flagging logic; falls back cleanly to
the runtime-only table when memory wasn't tracked. The workflow now
passes --track-memory so the comparison actually has memory data to
work with.
---
 .github/workflows/benchmark-regression.yml |   6 +-
 CHANGELOG.md                               |   7 +-
 tests/benchmarks/compare.py                | 122 +++++++++++++++++----
 tests/benchmarks/test_compare.py           |  84 ++++++++++++++
 4 files changed, 194 insertions(+), 25 deletions(-)

diff --git a/.github/workflows/benchmark-regression.yml b/.github/workflows/benchmark-regression.yml
index a26d17a..1fd2a55 100644
--- a/.github/workflows/benchmark-regression.yml
+++ b/.github/workflows/benchmark-regression.yml
@@ -7,7 +7,7 @@
 #     itself always comes from the PR checkout (the package uses a src/ layout,
 #     so the working tree never shadows the installed library).
 #   * Compare RELATIVE (%) change, not absolute seconds, so per-runner noise
-#     largely cancels.
+#     largely cancels. Both runtime and peak memory usage are compared.
 #   * Post the result as a single, updating ("sticky") PR comment. The check is
 #     informational and NON-BLOCKING for now; the threshold starts low and can
 #     be raised once the runner's real noise floor is known.
@@ -37,7 +37,9 @@ jobs:
     timeout-minutes: 20
     env:
       # Kept small so the whole job stays within a few minutes. Tune here.
-      BENCH_ARGS: "--min-N 50 --max-N 200 --num-steps 4 --archetype keisler"
+      # --track-memory is included so compare.py can also report peak-memory
+      # deltas alongside runtime (per the discussion on #144).
+      BENCH_ARGS: "--min-N 50 --max-N 200 --num-steps 4 --archetype keisler --track-memory"
       # Start low; raise once we've seen the runner's noise floor across a few
       # real runs (per the discussion on #144).
       THRESHOLD_PCT: "0.1"
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 59b063d..4f298d6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,9 +35,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add an automated benchmark regression check in CI: the scaling benchmark is
   run on the PR and its base branch back-to-back on the same runner (swapping
   only the library under test), and a sticky pull-request comment reports the
-  relative runtime change per grid size. Adds `tests/benchmarks/compare.py` and
-  a GitHub Actions workflow; informational and non-blocking for now, with a low
-  starting threshold to be calibrated against the runner's noise floor.
+  relative runtime and peak-memory change per grid size. Adds
+  `tests/benchmarks/compare.py` and a GitHub Actions workflow; informational
+  and non-blocking for now, with a low starting threshold to be calibrated
+  against the runner's noise floor.
   [\#144](https://github.com/mllam/weather-model-graphs/issues/144), @prajwal-tech07
 
 ### Deprecated
diff --git a/tests/benchmarks/compare.py b/tests/benchmarks/compare.py
index dc03684..7af331a 100644
--- a/tests/benchmarks/compare.py
+++ b/tests/benchmarks/compare.py
@@ -6,8 +6,10 @@
     [{"grid_points": int, "runtime_s": float, "peak_memory_mb": float | null}, ...]
 
 Given a *baseline* run (typically ``main``) and a *contender* run (the PR), it
-matches records by ``grid_points``, computes the relative change in runtime, and
-renders a Markdown table suitable for posting as a sticky pull-request comment.
+matches records by ``grid_points`` and computes the relative change in both
+runtime and peak memory usage (when both runs recorded memory, i.e. were run
+with ``--track-memory``), rendering a Markdown table suitable for posting as a
+sticky pull-request comment.
 
 The design follows the discussion in
 https://github.com/mllam/weather-model-graphs/issues/144: we compare relative
@@ -73,7 +75,13 @@ def _pct_change(baseline: float, contender: float) -> Optional[float]:
 
 
 class Row:
-    """A single grid-size comparison line in the report."""
+    """A single grid-size comparison line in the report.
+
+    Memory fields are ``None`` when either run didn't record
+    ``peak_memory_mb`` (i.e. wasn't run with ``--track-memory``) for this grid
+    size; the memory columns are omitted from the report entirely when no row
+    has memory data at all.
+    """
 
     def __init__(
         self,
@@ -82,12 +90,28 @@ def __init__(
         contender_s: float,
         delta_pct: Optional[float],
         is_regression: bool,
+        baseline_mem_mb: Optional[float] = None,
+        contender_mem_mb: Optional[float] = None,
+        mem_delta_pct: Optional[float] = None,
+        mem_is_regression: bool = False,
     ):
         self.grid_points = grid_points
         self.baseline_s = baseline_s
         self.contender_s = contender_s
         self.delta_pct = delta_pct
         self.is_regression = is_regression
+        self.baseline_mem_mb = baseline_mem_mb
+        self.contender_mem_mb = contender_mem_mb
+        self.mem_delta_pct = mem_delta_pct
+        self.mem_is_regression = mem_is_regression
+
+    @property
+    def has_memory(self) -> bool:
+        return self.baseline_mem_mb is not None and self.contender_mem_mb is not None
+
+    @property
+    def any_regression(self) -> bool:
+        return self.is_regression or self.mem_is_regression
 
 
 def compare(
@@ -95,20 +119,50 @@ def compare(
     contender: Dict[int, dict],
     threshold_pct: float,
 ) -> List[Row]:
-    """Build the per-grid-size comparison rows for the runtime metric.
+    """Build the per-grid-size comparison rows for runtime and peak memory.
 
     Only ``grid_points`` present in *both* runs are compared; unmatched sizes
     are skipped (they show up in the PR diff of the benchmark itself, so there is
     nothing to compare against). Rows are returned sorted by ``grid_points``.
+
+    The memory delta for a row is only computed when *both* the baseline and
+    contender records have a non-null ``peak_memory_mb`` (i.e. both runs used
+    ``--track-memory``); otherwise the row's memory fields stay ``None``.
     """
     common = sorted(set(baseline) & set(contender))
     rows: List[Row] = []
     for gp in common:
-        b = float(baseline[gp]["runtime_s"])
-        c = float(contender[gp]["runtime_s"])
+        b_record = baseline[gp]
+        c_record = contender[gp]
+
+        b = float(b_record["runtime_s"])
+        c = float(c_record["runtime_s"])
         delta = _pct_change(b, c)
         is_regression = delta is not None and delta > threshold_pct
-        rows.append(Row(gp, b, c, delta, is_regression))
+
+        b_mem = b_record.get("peak_memory_mb")
+        c_mem = c_record.get("peak_memory_mb")
+        mem_delta = None
+        mem_is_regression = False
+        if b_mem is not None and c_mem is not None:
+            b_mem = float(b_mem)
+            c_mem = float(c_mem)
+            mem_delta = _pct_change(b_mem, c_mem)
+            mem_is_regression = mem_delta is not None and mem_delta > threshold_pct
+
+        rows.append(
+            Row(
+                gp,
+                b,
+                c,
+                delta,
+                is_regression,
+                b_mem,
+                c_mem,
+                mem_delta,
+                mem_is_regression,
+            )
+        )
     return rows
 
 
@@ -126,6 +180,11 @@ def _fmt_delta(delta: Optional[float], is_regression: bool) -> str:
     return f"{delta:+.1f}% {icon}"
 
 
+def _fmt_mb(value: float) -> str:
+    """Format a peak-memory reading in MB."""
+    return f"{value:.1f}MB"
+
+
 def render_markdown(
     rows: List[Row],
     threshold_pct: float,
@@ -143,26 +202,49 @@ def render_markdown(
         )
         return "\n".join(lines) + "\n"
 
-    regressions = [r for r in rows if r.is_regression]
+    has_memory = any(r.has_memory for r in rows)
+    metric_label = "runtime/memory" if has_memory else "runtime"
+
+    regressions = [r for r in rows if r.any_regression]
     if regressions:
         lines.append(
             f"⚠️ **{len(regressions)} of {len(rows)}** grid sizes exceed the "
-            f"**+{threshold_pct:g}%** runtime threshold."
+            f"**+{threshold_pct:g}%** {metric_label} threshold."
         )
     else:
         lines.append(
-            f"✅ No runtime regression above **+{threshold_pct:g}%** "
+            f"✅ No {metric_label} regression above **+{threshold_pct:g}%** "
             f"across {len(rows)} grid sizes."
         )
     lines.append("")
 
-    lines.append(f"| grid points | {baseline_label} | {contender_label} | Δ runtime |")
-    lines.append("|---:|---:|---:|---:|")
+    header = ["grid points", baseline_label, contender_label, "Δ runtime"]
+    if has_memory:
+        header += [
+            f"{baseline_label} peak mem",
+            f"{contender_label} peak mem",
+            "Δ memory",
+        ]
+    lines.append("| " + " | ".join(header) + " |")
+    lines.append("|" + "|".join(["---:"] * len(header)) + "|")
+
     for r in rows:
-        lines.append(
-            f"| {r.grid_points:,} | {_fmt_seconds(r.baseline_s)} | "
-            f"{_fmt_seconds(r.contender_s)} | {_fmt_delta(r.delta_pct, r.is_regression)} |"
-        )
+        row_cells = [
+            f"{r.grid_points:,}",
+            _fmt_seconds(r.baseline_s),
+            _fmt_seconds(r.contender_s),
+            _fmt_delta(r.delta_pct, r.is_regression),
+        ]
+        if has_memory:
+            if r.has_memory:
+                row_cells += [
+                    _fmt_mb(r.baseline_mem_mb),
+                    _fmt_mb(r.contender_mem_mb),
+                    _fmt_delta(r.mem_delta_pct, r.mem_is_regression),
+                ]
+            else:
+                row_cells += ["n/a", "n/a", "n/a"]
+        lines.append("| " + " | ".join(row_cells) + " |")
 
     if unmatched:
         pretty = ", ".join(f"{gp:,}" for gp in unmatched)
@@ -208,9 +290,9 @@ def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
         "--threshold-pct",
         type=float,
         default=0.1,
-        help="Flag a grid size when its runtime increases by more than this "
-        "percentage. Start low and raise it once the runner's noise floor is "
-        "known (default: 0.1).",
+        help="Flag a grid size when its runtime or (if tracked) peak memory "
+        "increases by more than this percentage. Start low and raise it once "
+        "the runner's noise floor is known (default: 0.1).",
     )
     parser.add_argument(
         "--baseline-label", default="main", help="Column label for the baseline."
@@ -270,7 +352,7 @@ def main(argv: Optional[List[str]] = None) -> int:
         with open(args.output, "w", encoding="utf-8") as f:
             f.write(markdown)
 
-    if args.fail_on_regression and any(r.is_regression for r in rows):
+    if args.fail_on_regression and any(r.any_regression for r in rows):
         return 1
     return 0
 
diff --git a/tests/benchmarks/test_compare.py b/tests/benchmarks/test_compare.py
index b7852cd..d87e77b 100644
--- a/tests/benchmarks/test_compare.py
+++ b/tests/benchmarks/test_compare.py
@@ -91,6 +91,40 @@ def test_zero_baseline_row_is_not_a_regression():
     assert rows[0].is_regression is False
 
 
+def test_memory_delta_computed_when_both_sides_present():
+    baseline = {1024: _rec(1024, 1.0, peak_memory_mb=100.0)}
+    contender = {1024: _rec(1024, 1.0, peak_memory_mb=130.0)}  # +30%
+    rows = compare.compare(baseline, contender, threshold_pct=25.0)
+    row = rows[0]
+    assert row.has_memory is True
+    assert row.mem_delta_pct == pytest.approx(30.0)
+    assert row.mem_is_regression is True
+    assert row.any_regression is True  # driven by memory even if runtime is flat
+    assert row.is_regression is False  # runtime itself did not regress
+
+
+def test_memory_delta_none_when_either_side_missing():
+    # Contender never tracked memory (peak_memory_mb defaults to None).
+    baseline = {1024: _rec(1024, 1.0, peak_memory_mb=100.0)}
+    contender = {1024: _rec(1024, 1.0)}
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    row = rows[0]
+    assert row.has_memory is False
+    assert row.mem_delta_pct is None
+    assert row.mem_is_regression is False
+    assert row.any_regression is False
+
+
+def test_memory_regression_alone_does_not_affect_runtime_flag():
+    baseline = {1024: _rec(1024, 1.0, peak_memory_mb=100.0)}
+    contender = {1024: _rec(1024, 1.0, peak_memory_mb=200.0)}  # +100% memory
+    rows = compare.compare(baseline, contender, threshold_pct=1.0)
+    row = rows[0]
+    assert row.is_regression is False  # runtime unchanged
+    assert row.mem_is_regression is True
+    assert row.any_regression is True
+
+
 def test_load_results_rejects_empty(tmp_path):
     path = _write(tmp_path, "empty.json", [])
     with pytest.raises(ValueError, match="non-empty list"):
@@ -135,6 +169,47 @@ def test_render_markdown_no_overlap_message():
     assert "No overlapping grid sizes" in md
 
 
+def test_render_markdown_omits_memory_columns_when_absent():
+    rows = compare.compare(
+        {1024: _rec(1024, 1.0)}, {1024: _rec(1024, 1.0)}, threshold_pct=0.1
+    )
+    md = compare.render_markdown(rows, 0.1, "main", "PR")
+    assert "peak mem" not in md
+    assert "Δ memory" not in md
+
+
+def test_render_markdown_includes_memory_columns_when_present():
+    baseline = {1024: _rec(1024, 1.0, peak_memory_mb=100.0)}
+    contender = {1024: _rec(1024, 1.0, peak_memory_mb=130.0)}
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    md = compare.render_markdown(rows, 0.1, "main", "PR")
+    assert (
+        "| grid points | main | PR | Δ runtime | main peak mem | PR peak mem | Δ memory |"
+        in md
+    )
+    assert "100.0MB" in md
+    assert "130.0MB" in md
+    assert "+30.0% ⚠️" in md
+
+
+def test_render_markdown_shows_na_for_row_missing_memory():
+    # 1024 has memory on both sides, 4096 is missing it on the contender.
+    baseline = {
+        1024: _rec(1024, 1.0, peak_memory_mb=100.0),
+        4096: _rec(4096, 2.0, peak_memory_mb=200.0),
+    }
+    contender = {
+        1024: _rec(1024, 1.0, peak_memory_mb=110.0),
+        4096: _rec(4096, 2.0),  # no memory tracked this run
+    }
+    rows = compare.compare(baseline, contender, threshold_pct=0.1)
+    md = compare.render_markdown(rows, 0.1, "main", "PR")
+    assert "peak mem" in md  # column present because 1024 has data
+    lines = [line for line in md.splitlines() if line.startswith("| 4,096")]
+    assert len(lines) == 1
+    assert lines[0].count("n/a") == 3  # baseline/contender/delta memory cells
+
+
 def test_build_report_end_to_end(tmp_path):
     base = _write(tmp_path, "main.json", [_rec(1024, 1.0), _rec(4096, 2.0)])
     cont = _write(tmp_path, "pr.json", [_rec(1024, 1.0), _rec(4096, 3.0)])  # +50%
@@ -154,6 +229,15 @@ def test_main_fail_on_regression_exit_code(tmp_path, capsys):
     assert rc2 == 0
 
 
+def test_main_fail_on_regression_triggered_by_memory_alone(tmp_path):
+    # Runtime is flat; only peak memory regresses. --fail-on-regression must
+    # still catch it since any_regression considers both metrics.
+    base = _write(tmp_path, "main.json", [_rec(1024, 1.0, peak_memory_mb=100.0)])
+    cont = _write(tmp_path, "pr.json", [_rec(1024, 1.0, peak_memory_mb=200.0)])
+    rc = compare.main([base, cont, "--threshold-pct", "1.0", "--fail-on-regression"])
+    assert rc == 1
+
+
 def test_main_writes_output_file(tmp_path):
     base = _write(tmp_path, "main.json", [_rec(1024, 1.0)])
     cont = _write(tmp_path, "pr.json", [_rec(1024, 1.0)])