This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 8f3285c2 ci: shard the benchmark workflow across parallel jobs (#958)
8f3285c2 is described below
commit 8f3285c2ae42bf439af47a6bb2d531f8454b5e40
Author: Matt Topol <[email protected]>
AuthorDate: Sun Sep 20 16:49:16 2026 -0700
ci: shard the benchmark workflow across parallel jobs (#958)
### Rationale for this change
The `Benchmarks` workflow runs `go test -bench=. ./...` sequentially
across every
package. Because `go test`'s `-timeout` is applied per package, the
wall-clock time
is the sum over all packages — recent `main` runs have taken roughly 3.3
hours
(198–208 min).
### What changes are included in this PR?
Split the benchmark run so it can be parallelized, then combine the
results into a
single upload:
- **`ci/scripts/bench.sh`** — adds `--run` (benchmark a subset of
packages, writing
raw output to a `.dat`) and `--aggregate` (merge one or more `.dat`
files into a
single `bench_stats.json` via `gobenchdata`) modes. The existing
`bench.sh <dir> [--json|-json]` interface is unchanged, so nothing else
that calls
it needs to change.
- **`ci/scripts/bench_shard.sh`** (new) — prints a GitHub Actions matrix
that buckets
the packages containing benchmarks into N shards.
- **`.github/workflows/benchmark.yml`** — reworked into three jobs:
`setup` (compute
the shard matrix) → `benchmark` (matrix; each shard runs its packages
and uploads
its `.dat`) → `combine` (download all `.dat`, aggregate into one
`bench_stats.json`,
and — only on push to `main` — upload once to Conbench).
- **`ci/scripts/bench_adapt.py`** — reuses an existing
`bench_stats.json` (produced by
`combine`) instead of re-running the whole suite.
Because the shards are merged into one JSON and uploaded once, Conbench
still sees a
single run (no `run_id` fragmentation).
### Are these changes tested?
Locally:
- `shellcheck` clean on both scripts; `actionlint` clean on the
workflow; `py_compile`
OK on `bench_adapt.py`.
- Verified the split→merge end to end: ran `--run` on two packages, then
`--aggregate`
produced one `bench_stats.json` containing both suites, in the exact
shape
`bench_adapt.py` consumes.
- The legacy `bench.sh <dir> --json` path still runs → aggregates →
cleans up.
Opened as a **draft** to exercise the reworked workflow in CI end to end
(it triggers
on changes to these files).
### Are there any user-facing changes?
No. This only touches CI / benchmark tooling.
### Notes / follow-ups
- Sharding is currently round-robin by package, not runtime-weighted, so
a single
shard can hold two heavy packages (e.g. `arrow/compute` +
`parquet/internal/encoding`)
and become the long pole. The per-package `-timeout` (40m) remains the
hard floor for
any single package. Once this runs, per-shard timings can seed a
runtime-weighted
split or tune the shard count.
---------
Signed-off-by: Matt Topol <[email protected]>
---
.github/workflows/benchmark.yml | 86 ++++++++++--
arrow/array/json_reader_test.go | 9 +-
ci/scripts/bench.sh | 247 +++++++++++++++++++++++++++++++---
ci/scripts/bench_adapt.py | 291 ++++++++++++++++++++++++++++++++--------
ci/scripts/bench_shard.sh | 50 +++++++
5 files changed, 596 insertions(+), 87 deletions(-)
diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
index c40e0c2f..69f7b86e 100644
--- a/.github/workflows/benchmark.yml
+++ b/.github/workflows/benchmark.yml
@@ -23,36 +23,102 @@ on:
paths:
- ".github/workflows/benchmark.yml"
- "ci/scripts/bench.sh"
+ - "ci/scripts/bench_shard.sh"
- "ci/scripts/bench_adapt.py"
workflow_dispatch:
permissions:
contents: read
jobs:
+ setup:
+ runs-on: ubuntu-latest
+ outputs:
+ matrix: ${{ steps.shards.outputs.matrix }}
+ steps:
+ - name: Checkout repository
+ uses: actions/[email protected]
+ - name: Compute benchmark shards
+ id: shards
+ run: echo "matrix=$(bash ci/scripts/bench_shard.sh 6)" >>
"$GITHUB_OUTPUT"
benchmark:
+ needs: setup
runs-on: ubuntu-latest
+ # Shards run under a 90m watchdog (--timeout below); the slowest
+ # (./arrow/array) takes ~40m. This cap sits 10m above the watchdog so the
+ # watchdog trips first and names the stuck benchmark, ~4x sooner than the
+ # 6h GitHub default, while leaving room for checkout, setup and upload.
+ timeout-minutes: 100
strategy:
+ fail-fast: false
matrix:
- go: ['1.26.1']
- arch: ['amd64']
+ include: ${{ fromJson(needs.setup.outputs.matrix) }}
steps:
- name: Checkout repository
uses: actions/[email protected]
with:
submodules: recursive
- - name: Set up Python
+ - name: Install Go for Benchmarks
+ uses: actions/[email protected]
+ with:
+ go-version-file: go.mod
+ cache: true
+ cache-dependency-path: go.sum
+ check-latest: false
+ - name: Set up Python for benchmark metadata
uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 #
v7.0.0
with:
python-version: '3.9'
- - name: Install Go ${{ matrix.go }} for Benchmarks
+ - name: Run benchmark shard ${{ matrix.id }}
+ run: |
+ bash ci/scripts/bench.sh "$(pwd)" --run \
+ --packages "${{ matrix.packages }}" \
+ --out "bench_stat_${{ matrix.id }}.dat" \
+ --timeout 90m
+ - name: Install benchmark metadata collector
+ run: python3 -m pip install
benchadapt@git+https://github.com/conbench/conbench.git@3af4a55206ad3918762cc8dd7d3012eadbe96a54#subdirectory=benchadapt/python
+ - name: Capture benchmark machine information
+ env:
+ CONBENCH_MACHINE_INFO_NAME: amd64-debian-12
+ run: python3 ci/scripts/bench_adapt.py --capture-machine-info
"bench_stat_${{ matrix.id }}.dat"
+ - name: Upload shard results
+ uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a
# v7.0.1
+ with:
+ name: bench-dat-${{ matrix.id }}
+ path: |
+ bench_stat_${{ matrix.id }}.dat
+ bench_stat_${{ matrix.id }}.dat.machine.json
+ if-no-files-found: error
+ retention-days: 1
+ combine:
+ needs: benchmark
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout repository
+ uses: actions/[email protected]
+ - name: Install Go for Benchmarks
uses: actions/[email protected]
with:
- go-version: ${{ matrix.go }}
+ go-version-file: go.mod
cache: true
cache-dependency-path: go.sum
check-latest: false
- - name: Run Benchmarks
- if: github.event_name != 'push'
- run: bash ci/scripts/bench.sh $(pwd) --json
+ - name: Download shard results
+ uses:
actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
+ with:
+ pattern: bench-dat-*
+ path: bench-dat
+ merge-multiple: true
+ - name: Combine shard results
+ run: bash ci/scripts/bench.sh "$(pwd)" --aggregate --dat
"bench-dat/bench_stat_*.dat" --json-out bench_stats.json
+ - name: Upload combined results
+ uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a
# v7.0.1
+ with:
+ name: bench-stats-json
+ path: bench_stats.json
+ - name: Set up Python
+ if: github.event_name == 'push' && github.repository ==
'apache/arrow-go' && github.ref_name == 'main'
+ uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 #
v7.0.0
+ with:
+ python-version: '3.9'
- name: Upload results
if: github.event_name == 'push' && github.repository ==
'apache/arrow-go' && github.ref_name == 'main'
env:
@@ -60,7 +126,7 @@ jobs:
CONBENCH_EMAIL: ${{ secrets.CONBENCH_EMAIL }}
CONBENCH_PASSWORD: ${{ secrets.CONBENCH_PASS }}
CONBENCH_REF: ${{ github.ref_name }}
- CONBENCH_MACHINE_INFO_NAME: ${{ matrix.arch }}-debian-12
+ CONBENCH_MACHINE_INFO_NAME: amd64-debian-12
run: |
python3 -m pip install
benchadapt@git+https://github.com/conbench/conbench.git@3af4a55206ad3918762cc8dd7d3012eadbe96a54#subdirectory=benchadapt/python
- python3 ci/scripts/bench_adapt.py
+ python3 ci/scripts/bench_adapt.py --results bench_stats.json
--machine-info-dir bench-dat
diff --git a/arrow/array/json_reader_test.go b/arrow/array/json_reader_test.go
index 4254347e..fe8ee33a 100644
--- a/arrow/array/json_reader_test.go
+++ b/arrow/array/json_reader_test.go
@@ -245,6 +245,13 @@ func TestJSONReaderExponentialNotation(t *testing.T) {
}
}
+// benchMetadataPad keeps each generated record around 500 bytes of payload.
+// It must stay printable: JSON-escaping non-printable bytes (e.g. the NUL
+// bytes a zeroed []byte would produce) turns every byte into a \uXXXX escape,
+// and unescaping those is quadratic in goccy/go-json, which made these
+// benchmarks take hours.
+var benchMetadataPad = strings.Repeat("x", 500)
+
func generateJSONData(n int) []byte {
records := make([]map[string]any, n)
for i := range n {
@@ -253,7 +260,7 @@ func generateJSONData(n int) []byte {
"name": fmt.Sprintf("record_%d", i),
"value": float64(i) * 1.5,
"active": i%2 == 0,
- "metadata": fmt.Sprintf("metadata_%d_%s", i,
make([]byte, 500)),
+ "metadata": fmt.Sprintf("metadata_%d_%s", i,
benchMetadataPad),
}
}
diff --git a/ci/scripts/bench.sh b/ci/scripts/bench.sh
index a472b673..077dad81 100644
--- a/ci/scripts/bench.sh
+++ b/ci/scripts/bench.sh
@@ -17,37 +17,242 @@
# specific language governing permissions and limitations
# under the License.
-# this will output the benchmarks to STDOUT but if `-json` or `--json` is
passed
-# as the second argument, it will create a file "bench_stats.json"
-# in the directory this is called from containing a json representation
+# --run and --aggregate split a benchmark run so CI can shard packages across
+# parallel jobs and then combine the shards into one JSON for a single upload.
+# See usage() for the full interface.
set -exo pipefail
-# Validate input arguments
-if [ -z "$1" ]; then
- echo "Error: Missing source directory argument"
+GOBENCHDATA_VERSION="v1.3.1"
+
+usage() {
+ cat >&2 <<'EOF'
+Usage:
+ bench.sh <source_dir> [--json|-json]
+ Run every benchmark (BENCH_PACKAGES, default "./..."); with --json/-json
+ also aggregate into "bench_stats.json". Removes .dat files afterwards.
+ bench.sh <source_dir> --run [--packages "<patterns>"] [--out <file>]
[--timeout <dur>]
+ Run benchmarks only and write raw output to <file> (default
+ "bench_stat.dat"); leaves it in place for later aggregation.
+ bench.sh <source_dir> --aggregate [--dat "<glob>"] [--json-out <file>]
+ Combine .dat files (default "<source_dir>/bench_*.dat") into <file>
+ (default "bench_stats.json") via gobenchdata.
+
+Environment:
+ BENCH_PACKAGES Default packages for full/--run modes (default "./...").
+ BENCH_TIMEOUT Wall-clock limit for the whole `go test` invocation
+ (default "4h", sized for "./..."); "0" disables it. Nonzero
+ values require GNU timeout (timeout/gtimeout) or Python 3.
+EOF
+}
+
+run_with_python_timeout() {
+ local duration="$1"
+ shift
+ python3 - "${duration}" "$@" <<'PY'
+import os
+import signal
+import subprocess
+import sys
+
+duration = sys.argv[1]
+if not duration:
+ sys.exit("Error: benchmark timeout cannot be empty")
+multipliers = {"s": 1, "m": 60, "h": 60 * 60, "d": 24 * 60 * 60}
+suffix = duration[-1]
+if suffix in multipliers:
+ value = duration[:-1]
+ multiplier = multipliers[suffix]
+else:
+ value = duration
+ multiplier = 1
+
+try:
+ timeout_seconds = float(value) * multiplier
+except ValueError:
+ sys.exit(f"Error: invalid benchmark timeout: {duration!r}")
+if timeout_seconds <= 0:
+ sys.exit(f"Error: benchmark timeout must be greater than zero:
{duration!r}")
+
+process = subprocess.Popen(sys.argv[2:], start_new_session=True)
+
+
+def forward_signal(signum, _frame):
+ try:
+ os.killpg(process.pid, signum)
+ except ProcessLookupError:
+ pass
+
+
+for forwarded_signal in (
+ signal.SIGHUP,
+ signal.SIGINT,
+ signal.SIGQUIT,
+ signal.SIGTERM,
+):
+ signal.signal(forwarded_signal, forward_signal)
+
+try:
+ returncode = process.wait(timeout=timeout_seconds)
+except subprocess.TimeoutExpired:
+ forward_signal(signal.SIGQUIT, None)
+ try:
+ process.wait(timeout=60)
+ except subprocess.TimeoutExpired:
+ try:
+ os.killpg(process.pid, signal.SIGKILL)
+ except ProcessLookupError:
+ pass
+ process.wait()
+ sys.exit(124)
+
+if returncode < 0:
+ returncode = 128 - returncode
+sys.exit(returncode)
+PY
+}
+
+run_benchmarks() {
+ local source_dir="$1" packages="$2" out_file="$3" timeout="$4"
+
+ PARQUET_TEST_DATA="${source_dir}/parquet-testing/data"
+ export PARQUET_TEST_DATA
+
+ # `go test -timeout` does not cover benchmarks: the testing package stops its
+ # alarm before running them, so a runaway benchmark would otherwise burn the
+ # full 6h GitHub Actions job limit. Bound the whole run with a process
+ # watchdog instead, so a shard that goes pathological fails fast and visibly.
+ local runner=()
+ if [ "${timeout}" != "0" ]; then
+ local timeout_command=""
+ local candidate
+ for candidate in timeout gtimeout; do
+ if command -v "${candidate}" >/dev/null 2>&1 &&
+ "${candidate}" --signal=QUIT --kill-after=1m 0 true >/dev/null 2>&1;
then
+ timeout_command="${candidate}"
+ break
+ fi
+ done
+
+ if [ -n "${timeout_command}" ]; then
+ runner=("${timeout_command}" --signal=QUIT --kill-after=1m "${timeout}")
+ elif command -v python3 >/dev/null 2>&1; then
+ runner=(run_with_python_timeout "${timeout}")
+ else
+ echo "Error: no compatible watchdog is available for benchmark timeout
'${timeout}'." >&2
+ echo "Install GNU coreutils (for example, 'brew install coreutils' on
macOS) or Python 3, or set BENCH_TIMEOUT=0 / pass --timeout 0 to disable the
timeout guard." >&2
+ return 1
+ fi
+ fi
+
+ pushd "${source_dir}" >/dev/null
+ # shellcheck disable=SC2086 # intentional word-splitting of package patterns
+ "${runner[@]}" go test -bench=. -benchmem -run='^$' ${packages} | tee
"${out_file}"
+ popd >/dev/null
+}
+
+aggregate_results() {
+ local dat_glob="$1" json_out="$2"
+ local dat_files=()
+ local dat_file
+
+ go install "go.bobheadxi.dev/gobenchdata@${GOBENCHDATA_VERSION}"
+ PATH="$(go env GOPATH)/bin:$PATH"
+ export PATH
+
+ while IFS= read -r dat_file; do
+ dat_files[${#dat_files[@]}]="${dat_file}"
+ done < <(compgen -G "${dat_glob}")
+ if [ "${#dat_files[@]}" -eq 0 ]; then
+ echo "Error: no benchmark data files match: ${dat_glob}" >&2
+ return 1
+ fi
+ cat -- "${dat_files[@]}" | gobenchdata --json "${json_out}"
+}
+
+if [ -z "${1:-}" ]; then
+ echo "Error: Missing source directory argument" >&2
+ usage
exit 1
fi
source_dir="$1"
+shift
-PARQUET_TEST_DATA="${source_dir}/parquet-testing/data"
-export PARQUET_TEST_DATA
+mode="${1:-}"
-pushd "${source_dir}"
+packages="${BENCH_PACKAGES:-./...}"
+# Sized for the whole "./..." suite; CI passes a tighter --timeout per shard.
+timeout="${BENCH_TIMEOUT:-4h}"
-# lots of benchmarks, they can take a while
-# the timeout is for *ALL* benchmarks together,
-# not per benchmark
-go test -bench=. -benchmem -timeout 40m -run=^$ ./... | tee bench_stat.dat
+case "${mode}" in
+"" | -json | --json)
+ run_benchmarks "${source_dir}" "${packages}" "bench_stat.dat" "${timeout}"
+ if [[ "${mode}" == "-json" || "${mode}" == "--json" ]]; then
+ aggregate_results "${source_dir}/bench_*.dat" "bench_stats.json"
+ fi
+ rm "${source_dir}"/bench_*.dat
+ ;;
-popd
+--run)
+ shift
+ out_file="bench_stat.dat"
+ while [ "$#" -gt 0 ]; do
+ case "$1" in
+ --packages)
+ packages="$2"
+ shift 2
+ ;;
+ --out)
+ out_file="$2"
+ shift 2
+ ;;
+ --timeout)
+ timeout="$2"
+ shift 2
+ ;;
+ *)
+ echo "Error: unknown --run option: $1" >&2
+ usage
+ exit 1
+ ;;
+ esac
+ done
+ run_benchmarks "${source_dir}" "${packages}" "${out_file}" "${timeout}"
+ ;;
-if [[ "$2" = "-json" || "$2" = "--json" ]]; then
- go install go.bobheadxi.dev/[email protected]
- PATH=$(go env GOPATH)/bin:$PATH
- export PATH
- cat "${source_dir}"/bench_*.dat | gobenchdata --json bench_stats.json
-fi
+--aggregate)
+ shift
+ dat_glob="${source_dir}/bench_*.dat"
+ json_out="bench_stats.json"
+ while [ "$#" -gt 0 ]; do
+ case "$1" in
+ --dat)
+ dat_glob="$2"
+ shift 2
+ ;;
+ --json-out)
+ json_out="$2"
+ shift 2
+ ;;
+ *)
+ echo "Error: unknown --aggregate option: $1" >&2
+ usage
+ exit 1
+ ;;
+ esac
+ done
+ aggregate_results "${dat_glob}" "${json_out}"
+ ;;
+
+-h | --help)
+ usage
+ exit 0
+ ;;
-rm "${source_dir}"/bench_*.dat
+*)
+ echo "Error: unknown mode: ${mode}" >&2
+ usage
+ exit 1
+ ;;
+esac
diff --git a/ci/scripts/bench_adapt.py b/ci/scripts/bench_adapt.py
index 554538f4..f438eeeb 100644
--- a/ci/scripts/bench_adapt.py
+++ b/ci/scripts/bench_adapt.py
@@ -17,14 +17,16 @@
# specific language governing permissions and limitations
# under the License.
+import argparse
import json
+import logging
import os
import uuid
-import logging
from pathlib import Path
-from typing import List
+from typing import Dict, List, Optional, Tuple
from benchadapt import BenchmarkResult
+from benchadapt._machine_info import machine_info as collect_machine_info
from benchadapt.adapters import BenchmarkAdapter
from benchadapt.log import log
@@ -32,57 +34,196 @@ log.setLevel(logging.DEBUG)
ARROW_ROOT = Path(__file__).parent.parent.parent.resolve()
SCRIPTS_PATH = ARROW_ROOT / "ci" / "scripts"
+DEFAULT_RESULT_FILE = Path("bench_stats.json")
+MACHINE_INFO_PATTERN = "*.dat.machine.json"
+
-# `github_commit_info` is meant to communicate GitHub-flavored commit
-# information to Conbench. See
-#
https://github.com/conbench/conbench/blob/cf7931f/benchadapt/python/benchadapt/result.py#L66
-# for a specification.
-github_commit_info = {"repository": "https://github.com/apache/arrow-go"}
-
-if os.environ.get("CONBENCH_REF") == "main":
- # Assume GitHub Actions CI. The environment variable lookups below are
- # expected to fail when not running in GitHub Actions.
- github_commit_info = {
- "repository":
f'{os.environ["GITHUB_SERVER_URL"]}/{os.environ["GITHUB_REPOSITORY"]}',
- "commit": os.environ["GITHUB_SHA"],
- "pr_number": None, # implying default branch
+def upload_context() -> Tuple[Dict[str, object], str]:
+ # `github_commit_info` is meant to communicate GitHub-flavored commit
+ # information to Conbench. See
+ #
https://github.com/conbench/conbench/blob/cf7931f/benchadapt/python/benchadapt/result.py#L66
+ # for a specification.
+ github_commit_info: Dict[str, object] = {
+ "repository": "https://github.com/apache/arrow-go";
}
- run_reason = "commit"
-else:
- # Assume that the environment is not GitHub Actions CI. Error out if that
- # assumption seems to be wrong.
- assert os.getenv("GITHUB_ACTIONS") is None
-
- # This is probably a local dev environment, for testing. In this case, it
- # does usually not make sense to provide commit information (not a
- # controlled CI environment). Explicitly leave out "commit" and
"pr_number" to
- # reflect that (to not send commit information).
-
- # Reflect 'local dev' scenario in run_reason. Allow user to (optionally)
- # inject a custom piece of information into the run reason here, from
- # environment.
- run_reason = "localdev"
- custom_reason_suffix = os.getenv("CONBENCH_CUSTOM_RUN_REASON")
- if custom_reason_suffix is not None:
- run_reason += f" {custom_reason_suffix.strip()}"
+
+ if os.environ.get("CONBENCH_REF") == "main":
+ # Assume GitHub Actions CI. The environment variable lookups below are
+ # expected to fail when not running in GitHub Actions.
+ github_commit_info = {
+ "repository": (
+
f'{os.environ["GITHUB_SERVER_URL"]}/{os.environ["GITHUB_REPOSITORY"]}'
+ ),
+ "commit": os.environ["GITHUB_SHA"],
+ "pr_number": None, # implying default branch
+ }
+ run_reason = "commit"
+ else:
+ # Assume that the environment is not GitHub Actions CI. Error out if
+ # that assumption seems to be wrong.
+ assert os.getenv("GITHUB_ACTIONS") is None
+
+ # This is probably a local dev environment, for testing. In this case,
+ # it does usually not make sense to provide commit information (not a
+ # controlled CI environment). Explicitly leave out "commit" and
+ # "pr_number" to reflect that (to not send commit information).
+
+ # Reflect 'local dev' scenario in run_reason. Allow user to
(optionally)
+ # inject a custom piece of information into the run reason here, from
+ # environment.
+ run_reason = "localdev"
+ custom_reason_suffix = os.getenv("CONBENCH_CUSTOM_RUN_REASON")
+ if custom_reason_suffix is not None:
+ run_reason += f" {custom_reason_suffix.strip()}"
+
+ return github_commit_info, run_reason
+
+
+def capture_machine_info(dat_file: Path) -> Path:
+ packages = []
+ seen_packages = set()
+ with dat_file.open("r", encoding="utf-8") as source:
+ for line in source:
+ if not line.startswith("pkg: "):
+ continue
+ package = line.removeprefix("pkg: ").strip()
+ if not package:
+ raise ValueError(f"Empty 'pkg: ' header in {dat_file}")
+ if package not in seen_packages:
+ packages.append(package)
+ seen_packages.add(package)
+
+ if not packages:
+ raise ValueError(f"No 'pkg: ' headers found in {dat_file}")
+
+ sidecar = Path(f"{dat_file}.machine.json")
+ with sidecar.open("w", encoding="utf-8") as sink:
+ json.dump(
+ {
+ "packages": packages,
+ "machine_info": collect_machine_info(),
+ },
+ sink,
+ indent=2,
+ sort_keys=True,
+ )
+ sink.write("\n")
+ return sidecar
+
+
+def load_machine_info(directory: Path) -> Dict[str, Dict[str, object]]:
+ if not directory.is_dir():
+ raise ValueError(f"Machine metadata directory does not exist:
{directory}")
+
+ sidecars = sorted(directory.glob(MACHINE_INFO_PATTERN))
+ if not sidecars:
+ raise ValueError(
+ f"No machine metadata sidecars matching {MACHINE_INFO_PATTERN!r} "
+ f"found in {directory}"
+ )
+
+ machine_info_by_package: Dict[str, Dict[str, object]] = {}
+ source_by_package: Dict[str, Path] = {}
+ for sidecar in sidecars:
+ with sidecar.open("r", encoding="utf-8") as source:
+ payload = json.load(source)
+
+ if not isinstance(payload, dict) or set(payload) != {
+ "packages",
+ "machine_info",
+ }:
+ raise ValueError(
+ f"Invalid machine metadata schema in {sidecar}; expected only "
+ "'packages' and 'machine_info'"
+ )
+
+ packages = payload["packages"]
+ if (
+ not isinstance(packages, list)
+ or not packages
+ or not all(
+ isinstance(package, str)
+ and package
+ and package == package.strip()
+ for package in packages
+ )
+ ):
+ raise ValueError(f"Invalid package list in machine metadata
{sidecar}")
+ if len(packages) != len(set(packages)):
+ raise ValueError(f"Duplicate package in machine metadata
{sidecar}")
+
+ machine_info = payload["machine_info"]
+ if not isinstance(machine_info, dict) or not machine_info:
+ raise ValueError(f"Invalid machine_info in machine metadata
{sidecar}")
+
+ for package in packages:
+ if package in machine_info_by_package:
+ raise ValueError(
+ f"Ambiguous machine metadata for package {package!r}: "
+ f"{source_by_package[package]} and {sidecar}"
+ )
+ machine_info_by_package[package] = machine_info
+ source_by_package[package] = sidecar
+
+ return machine_info_by_package
class GoAdapter(BenchmarkAdapter):
- result_file = "bench_stats.json"
- command = ["bash", SCRIPTS_PATH / "bench.sh", ARROW_ROOT, "-json"]
+ def __init__(
+ self,
+ *args,
+ results_file: Optional[Path] = None,
+ machine_info_dir: Optional[Path] = None,
+ **kwargs,
+ ) -> None:
+ reuse_results = results_file is not None
+ if reuse_results != (machine_info_dir is not None):
+ raise ValueError(
+ "Explicit result reuse requires both results_file and
machine_info_dir"
+ )
- def __init__(self, *args, **kwargs) -> None:
- super().__init__(command=self.command, *args, **kwargs)
+ if reuse_results:
+ self.result_file = Path(results_file)
+ if not self.result_file.is_file():
+ raise ValueError(
+ f"Benchmark results file does not exist:
{self.result_file}"
+ )
+ self.machine_info_by_package =
load_machine_info(Path(machine_info_dir))
+ command = ["true"]
+ else:
+ # A pre-existing bench_stats.json must never opt the default path
+ # into reuse. The benchmark command overwrites it with a fresh run.
+ self.result_file = DEFAULT_RESULT_FILE
+ self.machine_info_by_package = None
+ command = ["bash", SCRIPTS_PATH / "bench.sh", ARROW_ROOT, "-json"]
+
+ self.github_commit_info, self.run_reason = upload_context()
+ super().__init__(command=command, *args, **kwargs)
def _transform_results(self) -> List[BenchmarkResult]:
- with open(self.result_file, "r") as f:
- raw_results = json.load(f)
+ with self.result_file.open("r", encoding="utf-8") as source:
+ raw_results = json.load(source)
+
+ suites = raw_results[0]["Suites"]
+ if self.machine_info_by_package is not None:
+ missing_packages = sorted(
+ {
+ suite["Pkg"]
+ for suite in suites
+ if suite["Pkg"] not in self.machine_info_by_package
+ }
+ )
+ if missing_packages:
+ raise ValueError(
+ "Missing machine metadata for benchmark package(s): "
+ + ", ".join(missing_packages)
+ )
run_id = uuid.uuid4().hex
parsed_results = []
- for suite in raw_results[0]["Suites"]:
+ for suite in suites:
batch_id = uuid.uuid4().hex
- pkg = suite["Pkg"]
+ package = suite["Pkg"]
for benchmark in suite["Benchmarks"]:
data = benchmark["Mem"]["MBPerSec"] * 1e6
@@ -92,39 +233,79 @@ class GoAdapter(BenchmarkAdapter):
ncpu = name[name.rfind("-") + 1 :]
pieces = name[: -(len(ncpu) + 1)].split("/")
- parsed = BenchmarkResult(
- run_id=run_id,
- batch_id=batch_id,
- stats={
+ result_fields = {
+ "run_id": run_id,
+ "batch_id": batch_id,
+ "stats": {
"data": [data],
"unit": "B/s",
"times": [time],
"time_unit": "i/s",
"iterations": benchmark["Runs"],
},
- context={
+ "context": {
"benchmark_language": "Go",
"goos": suite["Goos"],
"goarch": suite["Goarch"],
},
- tags={
- "pkg": pkg,
+ "tags": {
+ "pkg": package,
"num_cpu": ncpu,
"name": pieces[0],
"params": "/".join(pieces[1:]),
},
- run_reason=run_reason,
- github=github_commit_info,
- )
+ "run_reason": self.run_reason,
+ "github": self.github_commit_info,
+ }
+ if self.machine_info_by_package is not None:
+ result_fields["machine_info"] =
self.machine_info_by_package[package]
+
+ parsed = BenchmarkResult(**result_fields)
parsed.run_name = (
- f"{parsed.run_reason}: {github_commit_info.get('commit')}"
+ f"{parsed.run_reason}:
{self.github_commit_info.get('commit')}"
)
parsed_results.append(parsed)
return parsed_results
-if __name__ == "__main__":
- go_adapter = GoAdapter(result_fields_override={"info": {}})
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ mode = parser.add_mutually_exclusive_group()
+ mode.add_argument(
+ "--capture-machine-info",
+ type=Path,
+ metavar="DAT",
+ help="write DAT.machine.json using this machine and DAT's package
headers",
+ )
+ mode.add_argument(
+ "--results",
+ type=Path,
+ help="reuse an explicitly supplied gobenchdata JSON result file",
+ )
+ parser.add_argument(
+ "--machine-info-dir",
+ type=Path,
+ help="directory of per-shard .dat.machine.json provenance sidecars",
+ )
+ args = parser.parse_args()
+
+ if args.capture_machine_info is not None:
+ if args.machine_info_dir is not None:
+ parser.error("--machine-info-dir cannot be used with
--capture-machine-info")
+ capture_machine_info(args.capture_machine_info)
+ return
+
+ if (args.results is None) != (args.machine_info_dir is None):
+ parser.error("--results and --machine-info-dir must be supplied
together")
+
+ go_adapter = GoAdapter(
+ results_file=args.results,
+ machine_info_dir=args.machine_info_dir,
+ result_fields_override={"info": {}},
+ )
go_adapter()
-
\ No newline at end of file
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/ci/scripts/bench_shard.sh b/ci/scripts/bench_shard.sh
new file mode 100644
index 00000000..7a991757
--- /dev/null
+++ b/ci/scripts/bench_shard.sh
@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Print a GitHub Actions matrix (JSON) that splits the packages containing Go
+# benchmarks into up to <num_shards> groups, so bench.sh --run can execute them
+# in parallel jobs. Each element is {"id":N,"packages":"./pkg-a ./pkg-b ..."}.
+
+set -eo pipefail
+
+num_shards="${1:-6}"
+source_dir="${2:-.}"
+
+cd "${source_dir}"
+
+grep -rlE '^func Benchmark' --include='*_test.go' . |
+ xargs -n1 dirname |
+ sort -u |
+ awk -v want="${num_shards}" '
+ $0 == "" { next }
+ { pkg = ($0 ~ /^\.\//) ? $0 : "./" $0; pkgs[++count] = pkg }
+ END {
+ if (count == 0) { print "[]"; exit }
+ n = (want < count) ? want : count
+ for (k = 1; k <= count; k++) {
+ s = (k - 1) % n
+ shard[s] = shard[s] (shard[s] == "" ? "" : " ") pkgs[k]
+ }
+ printf "["
+ for (i = 0; i < n; i++) {
+ if (i > 0) printf ","
+ printf "{\"id\":%d,\"packages\":\"%s\"}", i, shard[i]
+ }
+ printf "]\n"
+ }'
