diff --git a/.github/workflows/benchmark-comment.yml b/.github/workflows/benchmark-comment.yml new file mode 100644 index 00000000..fe3b0b5a --- /dev/null +++ b/.github/workflows/benchmark-comment.yml @@ -0,0 +1,90 @@ +# 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. + +name: PR benchmark comment + +on: + workflow_run: + workflows: ['PR benchmark'] + types: [completed] + +permissions: + actions: read + pull-requests: write + +concurrency: + group: benchmark-comment-${{ github.event.workflow_run.head_repository.id }}-${{ github.event.workflow_run.head_branch }} + cancel-in-progress: false + +jobs: + comment: + if: ${{ github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.conclusion != 'cancelled' }} + runs-on: ubuntu-24.04 + timeout-minutes: 5 + steps: + # Never check out PR code, restore its caches, or execute artifact contents. + - name: Download this run's comment input + uses: actions/download-artifact@v4 + with: + name: benchmark-comment-${{ github.event.workflow_run.run_attempt }} + run-id: ${{ github.event.workflow_run.id }} + github-token: ${{ github.token }} + path: report + + - name: Create or update benchmark comment + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + const { owner, repo } = context.repo; + const run = context.payload.workflow_run; + if (run.path !== '.github/workflows/benchmark-pr.yml') { + throw new Error('Unexpected source workflow'); + } + // Artifact fields are untrusted: bind the requested PR to GitHub's run data. + const input = JSON.parse(fs.readFileSync('report/benchmark-context.json', 'utf8')); + if (!Number.isSafeInteger(input.number) || input.number <= 0) { + throw new Error('Invalid PR number'); + } + const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: input.number }); + if (pr.state !== 'open' || pr.base.repo.full_name !== `${owner}/${repo}` + || pr.head.repo?.id !== run.head_repository.id + || pr.head.ref !== run.head_branch || pr.head.sha !== run.head_sha + || input.head !== pr.head.sha || input.base !== pr.base.sha) { + core.info('Skipping unrelated or stale benchmark result'); + return; + } + // A delayed delivery from an earlier rerun must not replace the latest result. + const { data: current } = await github.rest.actions.getWorkflowRun({ owner, repo, run_id: run.id }); + if (current.run_attempt !== run.run_attempt) return; + const marker = ''; + const path = 'report/results/summary.md'; + const report = fs.existsSync(path) + ? fs.readFileSync(path, 'utf8') + : 'Benchmark did not produce a report. See the workflow logs.'; + if (Buffer.byteLength(report, 'utf8') > 60000) throw new Error('Report too large'); + const body = `${marker}\n[Benchmark run](${run.html_url}) · ${run.conclusion} · attempt ${run.run_attempt}\n\n${report}`; + const comments = await github.paginate(github.rest.issues.listComments, { + owner, repo, issue_number: pr.number, per_page: 100, + }); + const existing = comments.find(comment => + comment.user.login === 'github-actions[bot]' && comment.body?.startsWith(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: pr.number, body }); + } diff --git a/.github/workflows/benchmark-pr.yml b/.github/workflows/benchmark-pr.yml new file mode 100644 index 00000000..184e1c2d --- /dev/null +++ b/.github/workflows/benchmark-pr.yml @@ -0,0 +1,137 @@ +# 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. + +name: PR benchmark + +on: + pull_request: + branches: [main, 'release-*'] + paths: + - 'core/**' + - 'Cargo.toml' + - 'Cargo.lock' + - '.cargo/**' + - 'tools/benchmark_pr.py' + - 'tools/tests/test_benchmark_pr.py' + - 'tools/tests/test_benchmark_comment.cjs' + - '.github/workflows/benchmark-pr.yml' + - '.github/workflows/benchmark-comment.yml' + +permissions: + contents: read + +concurrency: + group: benchmark-pr-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + compare: + name: Compare base and PR + runs-on: ubuntu-24.04 + timeout-minutes: 25 + steps: + - name: Record comparison revisions + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + const pr = context.payload.pull_request; + fs.writeFileSync('benchmark-context.json', JSON.stringify({ + number: pr.number, head: pr.head.sha, base: pr.base.sha, + })); + + - name: Check out PR merge result + uses: actions/checkout@v4 + with: + ref: ${{ github.sha }} + path: candidate + persist-credentials: false + + - name: Check out exact base revision + uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.base.sha }} + path: base + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Set up stable Rust toolchain + run: | + rustup toolchain install stable --profile minimal + rustup default stable + env: + RUSTUP_TOOLCHAIN: stable + + - name: Cache Rust downloads + uses: actions/cache@v4 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + key: benchmark-${{ runner.os }}-rust-stable-${{ hashFiles('candidate/Cargo.lock', 'base/Cargo.lock') }} + restore-keys: | + benchmark-${{ runner.os }}-rust-stable- + + - name: Test comparison and failure handling + run: | + python3 -m unittest discover -s candidate/tools/tests -p 'test_benchmark_pr.py' -v + node candidate/tools/tests/test_benchmark_comment.cjs + + - name: Build and compare both revisions + # Stop the script before the job deadline, leaving time to publish failures. + run: python3 candidate/tools/benchmark_pr.py --base base --candidate candidate --output results --total-timeout 1200 + env: + RUSTUP_TOOLCHAIN: stable + + - name: Publish comparison summary + if: always() + run: | + if [ -f results/summary.md ]; then + cat results/summary.md >> "$GITHUB_STEP_SUMMARY" + else + echo 'Benchmark did not produce a report. Check the setup and build steps.' >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Upload report and raw samples + if: always() + uses: actions/upload-artifact@v4 + with: + name: pr-${{ github.event.pull_request.number }}-benchmark + overwrite: true + path: | + results/summary.md + results/summary.json + results/metadata.json + results/build-*.log + results/build-*.jsonl + results/raw/ + if-no-files-found: warn + retention-days: 30 + + - name: Upload comment input + if: always() + uses: actions/upload-artifact@v4 + with: + name: benchmark-comment-${{ github.run_attempt }} + path: | + benchmark-context.json + results/summary.md + retention-days: 7 diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs index a62566b1..5f685439 100644 --- a/core/benches/ann_bench.rs +++ b/core/benches/ann_bench.rs @@ -999,8 +999,51 @@ fn run_query_case( ("DISKANN", DiskAnnRawVectorEncoding::F16) => "f16", _ => "none", }; + // The original passes above retain their first-pass I/O and recall semantics. + // Optional CI timing starts after those complete sweeps have warmed each reader. + let steady_min_ms: u64 = read_env("ANN_STEADY_MIN_MS", 0)?; + if steady_min_ms > 0 && !storage.latency.is_zero() { + return Err("ANN_STEADY_MIN_MS requires local_ssd_warm_cache".into()); + } + let mut steady_sequential_queries = 0usize; + let mut steady_batch_queries = 0usize; + let mut steady_latencies = Vec::new(); + let mut steady_sequential_elapsed = Duration::ZERO; + let mut steady_batch_elapsed = Duration::ZERO; + if steady_min_ms > 0 { + let minimum = Duration::from_millis(steady_min_ms); + let started = Instant::now(); + while started.elapsed() < minimum { + for query in dataset.queries.chunks_exact(config.d) { + let query_started = Instant::now(); + std::hint::black_box(reader.search(query, search)?); + // Keep one complete warm sweep for P95; bound storage to nq samples. + if steady_sequential_queries == 0 { + steady_latencies.push(query_started.elapsed()); + } + } + steady_sequential_queries += config.nq; + } + steady_sequential_elapsed = started.elapsed(); + let started = Instant::now(); + while started.elapsed() < minimum { + std::hint::black_box(batch_reader.search_batch(&dataset.queries, config.nq, search)?); + steady_batch_queries += config.nq; + } + steady_batch_elapsed = started.elapsed(); + } + let steady_sequential_qps = if steady_sequential_queries == 0 { + 0.0 + } else { + steady_sequential_queries as f64 / steady_sequential_elapsed.as_secs_f64() + }; + let steady_batch_qps = if steady_batch_queries == 0 { + 0.0 + } else { + steady_batch_queries as f64 / steady_batch_elapsed.as_secs_f64() + }; println!( - "{dataset},{index},{storage},{n},{train_n},{raw_dataset_bytes},{nq},{d},{k},{nlist},{nprobe},{pq_m},{rq_bits},{build_distance},{raw_vector_encoding},{l_search},{build_ms},{train_ms},{add_ms},{write_ms},{peak_rss_bytes},{optimize_ms},{optimize_rounds},{optimize_ranges},{optimize_bytes},{file_bytes},{recall:.4},{first_us},{p50_us},{p95_us},{sequential_qps:.2},{seq_rounds},{seq_ranges},{seq_bytes},{batch_ms},{batch_qps:.2},{batch_rounds},{batch_ranges},{batch_bytes},{rq_seq_scanned},{rq_seq_refined},{rq_seq_final},{rq_seq_seeded_lists},{rq_seq_parallel_list_tasks},{rq_scanned},{rq_eligible},{rq_refined},{rq_refine_ratio:.6},{rq_final},{rq_refined_coarse_lookups},{rq_extra_plane_lookups},{rq_fastscan_blocks},{rq_scalar_blocks},{rq_seeded_lists},{rq_parallel_list_tasks}", + "{dataset},{index},{storage},{n},{train_n},{raw_dataset_bytes},{nq},{d},{k},{nlist},{nprobe},{pq_m},{rq_bits},{build_distance},{raw_vector_encoding},{l_search},{build_ms},{train_ms},{add_ms},{write_ms},{peak_rss_bytes},{optimize_ms},{optimize_rounds},{optimize_ranges},{optimize_bytes},{file_bytes},{recall:.4},{first_us},{p50_us},{p95_us},{sequential_qps:.2},{seq_rounds},{seq_ranges},{seq_bytes},{batch_ms},{batch_qps:.2},{batch_rounds},{batch_ranges},{batch_bytes},{rq_seq_scanned},{rq_seq_refined},{rq_seq_final},{rq_seq_seeded_lists},{rq_seq_parallel_list_tasks},{rq_scanned},{rq_eligible},{rq_refined},{rq_refine_ratio:.6},{rq_final},{rq_refined_coarse_lookups},{rq_extra_plane_lookups},{rq_fastscan_blocks},{rq_scalar_blocks},{rq_seeded_lists},{rq_parallel_list_tasks},{steady_min_ms},{steady_sequential_queries},{steady_batch_queries},{steady_sequential_ms},{steady_batch_ms},{steady_sequential_qps:.2},{steady_batch_qps:.2},{steady_sequential_p95_us:.3}", dataset = config.dataset_name, index = index.name, storage = storage.name, @@ -1059,6 +1102,9 @@ fn run_query_case( rq_scalar_blocks = rq_stats.scalar_blocks, rq_seeded_lists = rq_stats.seeded_lists, rq_parallel_list_tasks = rq_stats.parallel_list_tasks, + steady_sequential_ms = steady_sequential_elapsed.as_millis(), + steady_batch_ms = steady_batch_elapsed.as_millis(), + steady_sequential_p95_us = percentile(&steady_latencies, 95).as_secs_f64() * 1_000_000.0, ); Ok(()) } @@ -1067,7 +1113,7 @@ struct CsvRow; impl CsvRow { fn header() -> &'static str { - "dataset,index,storage,n,train_n,raw_dataset_bytes,nq,d,k,nlist,nprobe,pq_m,rq_bits,diskann_build_distance,diskann_raw_vector_encoding,l_search,build_ms,train_ms,add_ms,write_ms,peak_rss_bytes,optimize_ms,optimize_pread_rounds,optimize_pread_ranges,optimize_pread_bytes,file_bytes,recall_at_10,first_query_us,p50_query_us,p95_query_us,sequential_qps,sequential_pread_rounds,sequential_pread_ranges,sequential_pread_bytes,batch_ms,batch_qps,batch_pread_rounds,batch_pread_ranges,batch_pread_bytes,rq_sequential_scanned_vectors,rq_sequential_refined_vectors,rq_sequential_final_distance_evaluations,rq_sequential_seeded_lists,rq_sequential_parallel_list_tasks,rq_scanned_vectors,rq_eligible_vectors,rq_refined_vectors,rq_refine_ratio,rq_final_distance_evaluations,rq_refined_coarse_byte_lookups,rq_extra_plane_byte_lookups,rq_fastscan_blocks,rq_scalar_blocks,rq_seeded_lists,rq_parallel_list_tasks" + "dataset,index,storage,n,train_n,raw_dataset_bytes,nq,d,k,nlist,nprobe,pq_m,rq_bits,diskann_build_distance,diskann_raw_vector_encoding,l_search,build_ms,train_ms,add_ms,write_ms,peak_rss_bytes,optimize_ms,optimize_pread_rounds,optimize_pread_ranges,optimize_pread_bytes,file_bytes,recall_at_10,first_query_us,p50_query_us,p95_query_us,sequential_qps,sequential_pread_rounds,sequential_pread_ranges,sequential_pread_bytes,batch_ms,batch_qps,batch_pread_rounds,batch_pread_ranges,batch_pread_bytes,rq_sequential_scanned_vectors,rq_sequential_refined_vectors,rq_sequential_final_distance_evaluations,rq_sequential_seeded_lists,rq_sequential_parallel_list_tasks,rq_scanned_vectors,rq_eligible_vectors,rq_refined_vectors,rq_refine_ratio,rq_final_distance_evaluations,rq_refined_coarse_byte_lookups,rq_extra_plane_byte_lookups,rq_fastscan_blocks,rq_scalar_blocks,rq_seeded_lists,rq_parallel_list_tasks,steady_min_ms,steady_sequential_queries,steady_batch_queries,steady_sequential_ms,steady_batch_ms,steady_sequential_qps,steady_batch_qps,steady_sequential_p95_us" } } @@ -1096,7 +1142,9 @@ fn exact_ground_truth(dataset: &Dataset, k: usize) -> Vec> { .then_with(|| left.1.cmp(&right.1)) }); } - distances.into_iter().map(|(_, row)| row).collect() + // Borrow to collect a fresh top-k Vec. Consuming the iterator can + // reuse the much larger N-vector allocation via in-place collect. + distances.iter().map(|&(_, row)| row).collect() }) .collect() } diff --git a/tools/README.md b/tools/README.md index 17cac1cf..39375020 100644 --- a/tools/README.md +++ b/tools/README.md @@ -21,6 +21,89 @@ This directory contains helper scripts used by release managers and committers. +## PR / base benchmark + +The `PR benchmark` workflow compares the exact PR base SHA with GitHub's PR merge +commit on the same Ubuntu runner. It runs when the core, Cargo configuration, or +benchmark tooling changes. The result is +published in the Actions summary and posted as one updated bot comment, including +for external fork PRs. The comment opens with a five-row Recall/QPS table and +colored status markers. QPS drops greater than 10%/20% or recall drops greater +than 1/3 percentage points are yellow/red; these are advisory, not CI gates. +One collapsed table adds absolute QPS, P95, build time, process RSS and index size. + +A separate `workflow_run` publisher runs from the default branch with comment +permission. It downloads only the triggering run/attempt's report artifact, +checks the PR against GitHub's source repository, branch and revision data, and +skips closed PRs or changed head/base revisions. It never checks out PR code, +restores PR caches, or executes artifact contents. Raw CSVs, logs, environment +metadata and machine-readable results remain available as workflow artifacts. + +The publisher must first be merged into the target repository's default branch +before automatic comments can run. The PR introducing it still has its Actions +summary; subsequent runs can publish comments once the publisher is installed. + +`benchmark_pr.py` builds both revisions in release mode with separate target +directories, then runs six samples per version per index. Each index is tested +in a fresh process, alternating base/candidate and candidate/base pairs. Both +versions use the candidate's `ann_bench.rs` and its support module, so a change +to the benchmark itself cannot silently change the measurement between sides. +An incompatible driver/API combination fails explicitly and needs a compatible +shared driver before results can be compared. + +The initial workload is deliberately small: 10,000 synthetic 64D vectors, 4,096 +training vectors, 2,048 queries, top-10, seed 42, and two Rayon threads. It covers +IVF-FLAT, IVF-SQ, IVF-PQ, IVF-RQ and DiskANN on local warm page cache. First-pass +Recall and I/O preserve the original benchmark semantics: sequential queries +follow reader optimization and one first query; batch uses a separate optimized +reader. After those complete passes warm each reader, the timing phase repeats +full sequential and batch sweeps for at least one second each. The PR report's +QPS uses the full warm timing phase. P95 uses only its first complete sequential +sweep, bounding latency storage to one sample per query. First-pass metrics +remain in the CSV. +Warm timing is enabled by `ANN_STEADY_MIN_MS=1000` and is local-storage only. +Its elapsed times and query counts are included in the CSV for verification. + +This does not measure cold storage or real object-store performance. Recall is +currently measured on the first batch results. RSS is the process lifetime peak +up to build completion, including dataset and ground-truth allocations, not +index-only or search peak memory. Ground-truth IDs are copied from the retained +top-k slice to avoid retaining an N-vector allocation per query. + +The report shows medians, with relative QPS changes and recall changes in +percentage points. Sample ranges are kept in the JSON artifact. A recall +drop is highlighted alongside performance. This first version is informational: +performance/recall changes do not fail CI, but build failures, timeouts, invalid +metrics, missing samples, and mismatched workload parameters do. Sample ranges +are not confidence intervals, and a reported percentage is not a statistical +significance claim. When core sources and Cargo inputs are identical, the report +explicitly labels the run as A/A repeatability calibration. Such deltas cannot +demonstrate an index-code improvement. Repeated calibration runs are needed +before choosing any timing gate. + +To reproduce locally, create two **disposable** checkouts, use the same Rust +toolchain as the workflow, and run from the candidate checkout: + +```bash +python3 tools/benchmark_pr.py \ + --base /path/to/disposable-base \ + --candidate /path/to/disposable-candidate \ + --output /path/to/new-results-directory +``` + +The script overwrites the two benchmark-driver files in the base checkout. The +output directory must not exist. `--rounds 2` provides a shorter smoke run; +`--timeout` sets the per-sample timeout in seconds (default 180). +`--total-timeout` bounds all builds and samples together (default 1,200 seconds). +Each subprocess receives the smaller of its own timeout and the remaining total +budget. On timeout, its process group is killed and a failure report is written. +The 25-minute CI job reserves time beyond this 20-minute script budget for setup +and report upload. Dataset and +index options are pinned in the script; inherited `ANN_*` settings are removed. +CI uses the same Rust stable toolchain for both revisions and records the actual +compiler version in the report. It pins the x86-64 CPU target, caches dependency +downloads only, and uploads reports even when a comparison fails. + ## ANN-Benchmarks dataset conversion `convert_ann_benchmarks.py` converts a dense diff --git a/tools/benchmark_pr.py b/tools/benchmark_pr.py new file mode 100644 index 00000000..a052b2d2 --- /dev/null +++ b/tools/benchmark_pr.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +# 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. + +"""Compare two disposable checkouts using the candidate's ANN benchmark driver.""" + +import argparse +import csv +import hashlib +import json +import math +import os +from pathlib import Path +import platform +import shutil +import signal +import statistics +import subprocess +import time + + +INDEXES = ("IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "DISKANN") +WORKLOAD = { + "ANN_DATASET_NAME": "ci-clustered-v1", + "ANN_N": "10000", + "ANN_TRAIN_N": "4096", + "ANN_NQ": "2048", + "ANN_D": "64", + "ANN_K": "10", + "ANN_NLIST": "64", + "ANN_NPROBE": "8", + "ANN_PQ_M": "8", + "ANN_RQ_BITS": "4", + "ANN_CLUSTERS": "32", + "ANN_NOISE_DIMENSIONS": "64", + "ANN_SEED": "42", + "ANN_DISKANN_L_SEARCHES": "100", + "ANN_DISKANN_MEMORY_BUDGET_BYTES": str(256 * 1024 * 1024), + "ANN_DISKANN_BUILD_DISTANCE": "full_precision", + "ANN_DISKANN_RAW_VECTOR_ENCODING": "f32", + "ANN_STORAGE_CASES": "local_ssd_warm_cache", + "RAYON_NUM_THREADS": "2", + "ANN_STEADY_MIN_MS": "1000", +} +# These identify the workload, not measured results. Never compare unlike cases. +CASE_FIELDS = ( + "dataset", "index", "storage", "n", "train_n", "nq", "d", "k", "nlist", + "nprobe", "pq_m", "rq_bits", "diskann_build_distance", + "diskann_raw_vector_encoding", "l_search", "steady_min_ms", +) +METRICS = ( + ("recall_at_10", "Batch Recall@10", "recall"), + ("steady_sequential_qps", "Warm sequential QPS", "higher"), + ("steady_batch_qps", "Warm batch QPS", "higher"), + ("steady_sequential_p95_us", "Warm sequential P95 (µs)", "lower"), + ("build_ms", "Build (ms)", "lower"), + ("peak_rss_bytes", "Process peak RSS up to build completion (MiB)", "lower"), + ("file_bytes", "Index size (MiB)", "lower"), +) + + +def command_output(command, cwd=None): + return subprocess.check_output(command, cwd=cwd, text=True, timeout=30).strip() + + +def run_checked(command, *, deadline, timeout, **kwargs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("Benchmark total time budget exhausted; stopping before next process") + limit = min(timeout, remaining) + process = subprocess.Popen(command, start_new_session=os.name == "posix", **kwargs) + try: + returncode = process.wait(timeout=limit) + except subprocess.TimeoutExpired as error: + # Cargo can leave rustc children alive if only the parent is killed. + if os.name == "posix": + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + process.wait() + raise TimeoutError( + f"{command[0]} exceeded {limit:.1f}s process/remaining total budget" + ) from error + if returncode: + raise subprocess.CalledProcessError(returncode, command) + + +def numeric_field(row, field, path, *, integer=False): + raw = row.get(field) + if raw is None or not raw.strip(): + raise ValueError(f"{path}: missing or empty metric/parameter {field}") + try: + value = int(raw) if integer else float(raw) + except (ValueError, OverflowError) as error: + raise ValueError(f"{path}: invalid {field}={raw!r}; expected a number") from error + if not math.isfinite(value) or value < 0: + raise ValueError(f"{path}: invalid {field}={raw!r}; expected finite nonnegative value") + return value + + +def read_sample(path, index): + with path.open(newline="") as source: + rows = list(csv.DictReader(source)) + if len(rows) != 1: + raise ValueError(f"{path}: expected exactly one case, found {len(rows)}") + row = rows[0] + if row.get("index") != index: + raise ValueError(f"{path}: expected index {index}") + for field in CASE_FIELDS: + if not row.get(field): + raise ValueError(f"{path}: missing workload field {field}") + for field, _, _ in METRICS: + numeric_field(row, field, path) + if not 0 <= float(row["recall_at_10"]) <= 1 or numeric_field(row, "nq", path, integer=True) <= 0: + raise ValueError(f"{path}: invalid recall or query count") + if float(row["steady_sequential_qps"]) <= 0 or float(row["steady_batch_qps"]) <= 0: + raise ValueError(f"{path}: throughput must be positive") + minimum = numeric_field(row, "steady_min_ms", path, integer=True) + if minimum <= 0: + raise ValueError(f"{path}: steady_min_ms must be positive") + for mode in ("sequential", "batch"): + elapsed = numeric_field(row, f"steady_{mode}_ms", path) + queries = numeric_field(row, f"steady_{mode}_queries", path, integer=True) + if elapsed < minimum or queries < int(row["nq"]) or queries % int(row["nq"]): + raise ValueError(f"{path}: incomplete steady_{mode} measurement") + return row + + +def metric_value(row, field): + value = float(row[field]) + if field in ("peak_rss_bytes", "file_bytes"): + return value / (1024 * 1024) + return value + + +def delta_text(base, candidate, direction): + if direction == "recall": + return f"{(candidate - base) * 100:+.2f} pp" + if base == 0: + return "0.0%" if candidate == 0 else "n/a (base=0)" + return f"{(candidate / base - 1) * 100:+.1f}%" + + +def summarize(samples, rounds): + results = {} + for index in INDEXES: + sides = samples[index] + if any(len(sides[side]) != rounds for side in ("base", "candidate")): + raise ValueError(f"{index}: incomplete samples") + cases = { + tuple(row[field] for field in CASE_FIELDS) + for rows in sides.values() for row in rows + } + if len(cases) != 1: + raise ValueError(f"{index}: workload differs between samples") + result = {} + for field, label, direction in METRICS: + entry = {"label": label, "direction": direction} + for side in ("base", "candidate"): + values = [metric_value(row, field) for row in sides[side]] + entry[side] = { + "median": statistics.median(values), + "min": min(values), "max": max(values), "samples": values, + } + entry["delta"] = delta_text( + entry["base"]["median"], entry["candidate"]["median"], direction + ) + result[field] = entry + results[index] = result + return results + + +def query_alert_level(entry): + """Advisory levels for recall and QPS only; other metrics have no thresholds.""" + if entry["direction"] not in ("recall", "higher"): + raise ValueError("Query alerts only support recall and QPS") + base, candidate = entry["base"]["median"], entry["candidate"]["median"] + if entry["direction"] == "recall": + loss = round((base - candidate) * 100, 6) + yellow, red = 1, 3 + else: + loss = round((1 - candidate / base) * 100, 6) if base else 0 + yellow, red = 10, 20 + return 2 if loss > red else 1 if loss > yellow else 0 + + +def render_report(metadata, results): + focus = (("recall_at_10", "Recall"), ("steady_sequential_qps", "single QPS"), + ("steady_batch_qps", "batch QPS")) + levels = {index: max(query_alert_level(metrics[field]) for field, _ in focus) + for index, metrics in results.items()} + flagged = sum(level > 0 for level in levels.values()) + icon = ("🟢", "🟡", "🔴")[max(levels.values(), default=0)] + headline = (f"{flagged}/{len(results)} indexes need a look" if flagged + else f"No alerts ({len(results)}/{len(results)} indexes)") + lines = ["## Vector index benchmark", "", f"**{icon} {headline}**", ""] + if metadata.get("calibration"): + lines += ["A/A calibration — identical core code; deltas show measurement variation.", ""] + lines += ["| Index | Recall@10 (Δ) | Single QPS Δ | Batch QPS Δ | Status |", + "|---|---:|---:|---:|---|"] + for index, metrics in results.items(): + recall = metrics["recall_at_10"] + reasons = [label for field, label in focus if query_alert_level(metrics[field])] + status = ("🟢 OK" if not reasons else + f"{('🟢', '🟡', '🔴')[levels[index]]} {', '.join(reasons)}") + lines.append( + f"| {index} | {recall['candidate']['median'] * 100:.2f}% ({recall['delta']}) | " + f"{metrics['steady_sequential_qps']['delta']} | " + f"{metrics['steady_batch_qps']['delta']} | {status} |" + ) + lines += ["", "🟡 QPS ↓ >10% or Recall ↓ >1pp · 🔴 QPS ↓ >20% or Recall ↓ >3pp. " + "Advisory only; QPS is measured after warmup.", "", + "
Absolute metrics & environment", ""] + lines += render_details(metadata, results) + lines += ["", "
", ""] + return "\n".join(lines) + + +def render_details(metadata, results): + fields = ("steady_sequential_qps", "steady_batch_qps", "steady_sequential_p95_us", + "build_ms", "peak_rss_bytes", "file_bytes") + lines = [ + "Medians, base → PR. QPS is timed after warmup; P95 uses the first warm sweep.", "", + "| Index | Single QPS | Batch QPS | P95 (µs) | Build (ms) | Process RSS (MiB) | Size (MiB) |", + "|---|---:|---:|---:|---:|---:|---:|", + ] + for index, metrics in results.items(): + values = [f"{metrics[field]['base']['median']:,.2f} → " + f"{metrics[field]['candidate']['median']:,.2f}" for field in fields] + lines.append(f"| {index} | " + " | ".join(values) + " |") + lines += [ + "", + f"Base `{metadata['base_sha']}` · PR merge `{metadata['candidate_sha']}`", + f"{metadata['rustc'].splitlines()[0]} · {metadata['platform']} · {metadata['cpu']}", + f"{metadata['rounds']} samples/version/index, alternating execution order; " + "2 threads; 10k × 64D vectors; 2,048 queries; ≥1s per timed mode.", + "RSS is the process lifetime peak up to build completion, including dataset/ground truth.", + "Raw CSVs, sample ranges, workload settings and logs are in the workflow artifact.", + ] + return lines + + +def build(checkout, output, side, env, deadline): + # Separate target directories prevent artifacts from one revision leaking into the other. + command = ["cargo", "bench", "--locked", "-p", "paimon-vindex-core", "--bench", + "ann_bench", "--no-run", "--message-format=json", "--target-dir", + str(output / "build" / side)] + print(f"Building {side}", flush=True) + messages_path = output / f"build-{side}.jsonl" + with messages_path.open("w") as messages, (output / f"build-{side}.log").open("w") as log: + run_checked(command, cwd=checkout, env=env, stdout=messages, stderr=log, + deadline=deadline, timeout=900) + executables = [] + for line in messages_path.read_text().splitlines(): + message = json.loads(line) + if (message.get("reason") == "compiler-artifact" + and message.get("target", {}).get("name") == "ann_bench" + and message.get("executable")): + executables.append(message["executable"]) + if len(executables) != 1: + raise ValueError(f"{side}: expected one benchmark executable") + return executables[0] + + +def library_fingerprint(checkout): + paths = list((checkout / "core/src").rglob("*")) + paths += [checkout / name for name in ("Cargo.toml", "Cargo.lock", "core/Cargo.toml", "core/build.rs")] + paths += list((checkout / ".cargo").rglob("*")) + digest = hashlib.sha256() + for path in sorted(path for path in paths if path.is_file()): + digest.update(str(path.relative_to(checkout)).encode() + b"\0" + path.read_bytes()) + return digest.hexdigest() + + +def run(args): + deadline = time.monotonic() + args.total_timeout + output = args.output.resolve() + base, candidate = args.base.resolve(), args.candidate.resolve() + if base == candidate: + raise ValueError("base and candidate must be separate disposable checkouts") + if command_output(["git", "status", "--porcelain"], base): + raise ValueError("base checkout must be clean before copying the shared driver") + metadata = { + "base_sha": command_output(["git", "rev-parse", "HEAD"], base), + "candidate_sha": command_output(["git", "rev-parse", "HEAD"], candidate), + "rustc": command_output(["rustc", "--version", "--verbose"]), + "platform": platform.platform(), + "cpu": platform.processor(), + "rounds": args.rounds, "workload": WORKLOAD, + "candidate_dirty": bool(command_output(["git", "status", "--porcelain"], candidate)), + "execution_order": [], + "total_timeout_seconds": args.total_timeout, + "base_library_sha256": library_fingerprint(base), + "candidate_library_sha256": library_fingerprint(candidate), + "rustflags": "-C target-cpu=x86-64" if platform.machine() == "x86_64" else "", + } + metadata["calibration"] = metadata["base_library_sha256"] == metadata["candidate_library_sha256"] + if shutil.which("lscpu"): + metadata["cpu"] = next((line.split(":", 1)[1].strip() + for line in command_output(["lscpu"]).splitlines() + if line.startswith("Model name:")), metadata["cpu"]) + # Use candidate driver on both versions, including its support module. API incompatibility + # fails the run explicitly instead of silently comparing different benchmark programs. + digest = hashlib.sha256() + for relative in ("core/benches/ann_bench.rs", "core/benches/support/ann_bench_support.rs"): + data = (candidate / relative).read_bytes() + digest.update(relative.encode() + b"\0" + data) + (base / relative).parent.mkdir(parents=True, exist_ok=True) + (base / relative).write_bytes(data) + metadata["driver_sha256"] = digest.hexdigest() + (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + env = {key: value for key, value in os.environ.items() + if not key.startswith(("ANN_", "DISKANN_BENCH_", "CARGO_PROFILE_")) + and key not in ("RUSTFLAGS", "CARGO_ENCODED_RUSTFLAGS", "CARGO_BUILD_RUSTFLAGS")} + env.update(WORKLOAD) + env.update({"RUSTFLAGS": metadata["rustflags"], "LC_ALL": "C", "CARGO_INCREMENTAL": "0"}) + executables = {side: build(checkout, output, side, env, deadline) + for side, checkout in (("base", base), ("candidate", candidate))} + samples = {index: {"base": [], "candidate": []} for index in INDEXES} + raw = output / "raw" + raw.mkdir() + for index in INDEXES: + for repetition in range(args.rounds): + order = ("base", "candidate") if repetition % 2 == 0 else ("candidate", "base") + for side in order: + prefix = raw / f"{index}-{repetition + 1}-{side}" + print(f"Running {prefix.name}", flush=True) + sample_env = dict(env, ANN_INDEXES=index, ANN_OUTPUT_DIR=str(output / "indexes")) + started = time.monotonic() + with prefix.with_suffix(".csv").open("w") as csv_file, prefix.with_suffix(".log").open("w") as log: + run_checked([executables[side]], env=sample_env, cwd=output, + stdout=csv_file, stderr=log, timeout=args.timeout, deadline=deadline) + samples[index][side].append(read_sample(prefix.with_suffix(".csv"), index)) + metadata["execution_order"].append({ + "index": index, "round": repetition + 1, "side": side, + "wall_seconds": round(time.monotonic() - started, 3), + }) + (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + results = summarize(samples, args.rounds) + (output / "summary.json").write_text(json.dumps(results, indent=2) + "\n") + (output / "summary.md").write_text(render_report(metadata, results)) + print(f"Report: {output / 'summary.md'}", flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", type=Path, required=True, help="Disposable base checkout (driver is overwritten)") + parser.add_argument("--candidate", type=Path, required=True, help="Candidate checkout providing the shared driver") + parser.add_argument("--output", type=Path, required=True, help="New output directory") + parser.add_argument("--rounds", type=int, default=6, help="Samples per version/index; even number >= 2") + parser.add_argument("--timeout", type=int, default=180, help="Timeout in seconds per sample") + parser.add_argument("--total-timeout", type=int, default=1200, help="Total build/sample budget in seconds") + args = parser.parse_args() + if args.rounds < 2 or args.rounds % 2 or args.timeout <= 0 or args.total_timeout <= 0: + parser.error("rounds must be even and >= 2; timeouts must be positive") + args.output = args.output.resolve() + args.output.mkdir(parents=True, exist_ok=False) + try: + run(args) + except Exception as error: + (args.output / "summary.md").write_text( + "# PR / base benchmark failed\n\n" + "The comparison is incomplete; no performance conclusion is available. " + "See the uploaded build/sample logs for details.\n\n" + f"```text\n{error}\n```\n" + ) + raise + + +if __name__ == "__main__": + main() diff --git a/tools/tests/test_benchmark_comment.cjs b/tools/tests/test_benchmark_comment.cjs new file mode 100644 index 00000000..baf8352f --- /dev/null +++ b/tools/tests/test_benchmark_comment.cjs @@ -0,0 +1,76 @@ +// 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. + +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const workflow = fs.readFileSync(path.join(__dirname, '../../.github/workflows/benchmark-comment.yml'), 'utf8'); +const script = workflow.split(' script: |\n')[1].replace(/^ /gm, ''); +const execute = new (Object.getPrototypeOf(async function () {}).constructor)( + 'require', 'github', 'context', 'core', script); + +async function check({ mutate = () => {}, existing = false, missing = false, posts = 0 } = {}) { + const run = { id: 1, run_attempt: 1, path: '.github/workflows/benchmark-pr.yml', + head_repository: { id: 2 }, head_branch: 'feature', head_sha: 'head', + conclusion: 'success', html_url: 'https://github.com/owner/repo/actions/runs/1' }; + // Distinct repository IDs represent a fork PR; no workflow_run.pull_requests needed. + const pr = { number: 95, state: 'open', base: { sha: 'base', repo: { full_name: 'owner/repo' } }, + head: { sha: 'head', ref: 'feature', repo: { id: 2 } } }; + const input = { number: 95, head: 'head', base: 'base' }; + const current = { run_attempt: 1 }; + mutate({ run, pr, input, current }); + const calls = []; + const github = { rest: { + pulls: { get: async () => ({ data: pr }) }, + actions: { getWorkflowRun: async () => ({ data: current }) }, + issues: { listComments() {}, + updateComment: async value => calls.push(['update', value]), + createComment: async value => calls.push(['create', value]) }, + }, paginate: async () => existing ? [{ id: 7, user: { login: 'github-actions[bot]' }, + body: 'old' }] : [] }; + const files = { existsSync: () => !missing, readFileSync: name => { + if (name === 'report/benchmark-context.json') return JSON.stringify(input); + assert.equal(name, 'report/results/summary.md'); + return '🟢 Benchmark report'; + } }; + await execute(name => { assert.equal(name, 'fs'); return files; }, github, + { repo: { owner: 'owner', repo: 'repo' }, payload: { workflow_run: run } }, { info() {} }); + assert.equal(calls.length, posts); + if (posts) { + assert.equal(calls[0][0], existing ? 'update' : 'create'); + assert.match(calls[0][1].body, missing ? /did not produce/ : /🟢 Benchmark report/); + } +} + +(async () => { + await check({ posts: 1 }); + await check({ existing: true, posts: 1 }); + await check({ missing: true, posts: 1 }); + for (const mutate of [ + ({ pr }) => pr.state = 'closed', + ({ pr }) => pr.head.sha = 'new', + ({ pr }) => pr.base.sha = 'new', + ({ pr }) => pr.head.repo.id = 3, + ({ pr }) => pr.head.ref = 'other', + ({ pr }) => pr.base.repo.full_name = 'other/repo', + ({ input }) => input.head = 'other', + ({ current }) => current.run_attempt = 2, + ]) await check({ mutate }); + await assert.rejects(check({ mutate: ({ input }) => input.number = '../95' }), /Invalid PR/); + await assert.rejects(check({ mutate: ({ run }) => run.path = 'other.yml' }), /Unexpected source/); + console.log('Publisher checks passed: fork routing, create/update, missing report and stale/unrelated runs'); +})().catch(error => { console.error(error); process.exitCode = 1; }); diff --git a/tools/tests/test_benchmark_pr.py b/tools/tests/test_benchmark_pr.py new file mode 100644 index 00000000..86df1813 --- /dev/null +++ b/tools/tests/test_benchmark_pr.py @@ -0,0 +1,168 @@ +# 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. + +import copy +import csv +import importlib.util +from pathlib import Path +import tempfile +import subprocess +import sys +import time +import unittest + + +SPEC = importlib.util.spec_from_file_location( + "benchmark_pr", Path(__file__).resolve().parents[1] / "benchmark_pr.py" +) +bench = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(bench) + + +def row(index="IVF_FLAT"): + values = dict.fromkeys(bench.CASE_FIELDS, "1") + values.update({field: "100" for field, _, _ in bench.METRICS}) + values.update(index=index, nq="10", recall_at_10="0.95", + steady_sequential_ms="100", steady_batch_ms="100", + steady_sequential_queries="100", steady_batch_queries="100") + return values + + +def samples(): + return {index: {side: [row(index), row(index)] for side in ("base", "candidate")} + for index in bench.INDEXES} + + +class BenchmarkComparisonTest(unittest.TestCase): + def test_medians_units_and_recall_percentage_points(self): + values = samples() + values["IVF_FLAT"]["base"][1]["steady_sequential_qps"] = "300" + values["IVF_FLAT"]["candidate"][0]["steady_sequential_qps"] = "200" + values["IVF_FLAT"]["candidate"][1]["steady_sequential_qps"] = "400" + for value in values["IVF_FLAT"]["candidate"]: + value["recall_at_10"] = "0.90" + result = bench.summarize(values, 2)["IVF_FLAT"] + self.assertEqual(result["steady_sequential_qps"]["base"]["median"], 200) + self.assertEqual(result["steady_sequential_qps"]["delta"], "+50.0%") + self.assertEqual(result["recall_at_10"]["delta"], "-5.00 pp") + self.assertAlmostEqual(result["file_bytes"]["base"]["median"], 100 / (1024 * 1024)) + + def test_incomplete_or_mismatched_workloads_fail(self): + for change in ("missing", "shape", "parameters"): + with self.subTest(change=change): + values = samples() + if change == "missing": + values["IVF_FLAT"]["candidate"].pop() + else: + field = "nq" if change == "shape" else "nprobe" + values["IVF_FLAT"]["candidate"][0][field] = "999" + with self.assertRaises(ValueError): + bench.summarize(values, 2) + + def test_csv_rejects_missing_duplicate_and_invalid_results(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "sample.csv" + good = row() + bad_values = [] + for field, value in (("steady_batch_qps", "nan"), ("build_ms", "-1"), + ("recall_at_10", "1.1"), ("nq", "0"), + ("steady_batch_qps", "0"), ("index", "DISKANN")): + bad = copy.copy(good) + bad[field] = value + bad_values.append([bad]) + for rows in ([], [good, good], *bad_values): + with self.subTest(rows=rows): + with path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=good) + writer.writeheader() + writer.writerows(rows) + with self.assertRaises(ValueError): + bench.read_sample(path, "IVF_FLAT") + with path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=good) + writer.writeheader() + writer.writerow(good) + self.assertEqual(bench.read_sample(path, "IVF_FLAT"), good) + + def test_parse_errors_identify_file_field_and_value(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "broken.csv" + for field, value in (("steady_batch_qps", None), ("steady_batch_qps", ""), + ("steady_batch_qps", "oops"), ("nq", "2.5")): + with self.subTest(field=field, value=value): + sample = row() + if value is None: + del sample[field] + else: + sample[field] = value + with path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=sample) + writer.writeheader() + writer.writerow(sample) + with self.assertRaises(ValueError) as caught: + bench.read_sample(path, "IVF_FLAT") + self.assertIn(str(path), str(caught.exception)) + self.assertIn(field, str(caught.exception)) + if value: + self.assertIn(value, str(caught.exception)) + + def test_total_budget_stops_before_start_and_limits_running_process(self): + with tempfile.TemporaryDirectory() as directory: + marker = Path(directory) / "must-not-exist" + with self.assertRaisesRegex(TimeoutError, "total time budget"): + bench.run_checked([sys.executable, "-c", + "from pathlib import Path; Path(__import__('sys').argv[1]).touch()", + str(marker)], deadline=time.monotonic() - 1, timeout=10) + self.assertFalse(marker.exists()) + started = time.monotonic() + with self.assertRaisesRegex(TimeoutError, "remaining total budget"): + bench.run_checked([sys.executable, "-c", "import time; time.sleep(10)"], + deadline=started + 0.1, timeout=10, + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + self.assertLess(time.monotonic() - started, 3) + + def test_short_or_partial_steady_measurements_fail(self): + for field, value in (("steady_batch_ms", "0"), ("steady_sequential_queries", "11")): + with self.subTest(field=field), tempfile.TemporaryDirectory() as directory: + sample = row() + sample[field] = value + path = Path(directory) / "incomplete.csv" + with path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=sample) + writer.writeheader() + writer.writerow(sample) + with self.assertRaisesRegex(ValueError, "incomplete steady_"): + bench.read_sample(path, "IVF_FLAT") + + def test_zero_baseline_is_not_a_false_percentage(self): + self.assertEqual(bench.delta_text(0, 1, "lower"), "n/a (base=0)") + self.assertEqual(bench.delta_text(0, 0, "lower"), "0.0%") + + def test_loose_alert_boundaries(self): + for candidate, expected in ((110, 0), (90, 0), (89, 1), (80, 1), (79, 2)): + entry = dict(direction="higher", base={"median": 100}, candidate={"median": candidate}) + self.assertEqual(bench.query_alert_level(entry), expected) + for candidate, expected in ((0.94, 0), (0.939, 1), (0.92, 1), (0.919, 2)): + entry = dict(direction="recall", base={"median": 0.95}, candidate={"median": candidate}) + self.assertEqual(bench.query_alert_level(entry), expected) + + with self.assertRaises(ValueError): + bench.query_alert_level(dict(direction="lower", base={"median": 100}, candidate={"median": 120})) + + +if __name__ == "__main__": + unittest.main()