Skip to content

Use RDS for baseline #9093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AdamGS wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Use RDS for baseline #9093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 26 additions & 10 deletions .github/workflows/bench-pr.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ jobs:
VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1"
FLAT_LAYOUT_INLINE_ARRAY_NODE: "1"
run: |
python3 scripts/random-access-split.py
python3 scripts/random-access-split.py --emit-ingest-records

- name: Run ${{ matrix.benchmark.name }} benchmark
if: matrix.benchmark.id != 'random-access-bench'
Expand All @@ -104,30 +104,46 @@ jobs:
VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1"
FLAT_LAYOUT_INLINE_ARRAY_NODE: "1"
run: |
bash scripts/bench-taskset.sh target/release_debug/${{ matrix.benchmark.id }} -d gh-json -o results.json
bash scripts/bench-taskset.sh target/release_debug/${{ matrix.benchmark.id }} \
-d gh-json \
-o results.json \
--ingest-jsonl results.ingest.jsonl

- name: Setup AWS CLI
- name: Configure AWS credentials for RDS baseline (OIDC)
if: github.event.pull_request.head.repo.fork == false
uses: aws-actions/configure-aws-credentials@e6de054238d6b7531b4efff3b6587d9aade6a06c # v6
with:
role-to-assume: arn:aws:iam::245040174862:role/GitHubBenchmarkRole
aws-region: us-east-1
role-to-assume: ${{ vars.GH_BENCH_INGEST_ROLE_ARN }}
aws-region: ${{ vars.RDS_BENCH_REGION }}

- name: Install uv
uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6
with:
sync: false

- name: Compare results
if: github.event.pull_request.head.repo.fork == false
shell: bash
env:
RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }}
RDS_BENCH_DB_NAME: ${{ vars.RDS_BENCH_DB_NAME }}
AWS_REGION: ${{ vars.RDS_BENCH_REGION }}
run: |
set -Eeu -o pipefail -x

python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request
gzip -d -c data.json.gz > base.json

uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.benchmark.name }}" \
> comment.md
curl -fsSL https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem \
-o "${RUNNER_TEMP}/rds-global-bundle.pem"
DSN="postgresql://bench_ingest@${RDS_BENCH_INSTANCE_ENDPOINT}:5432/${RDS_BENCH_DB_NAME}?sslmode=verify-full&sslrootcert=${RUNNER_TEMP}/rds-global-bundle.pem"
uv run --no-project scripts/fetch-benchmark-baseline.py results.ingest.jsonl \
--postgres "${DSN}" \
--region "${AWS_REGION}" \
--output base.ingest.jsonl
uv run --no-project scripts/compare-benchmark-jsons.py \
--ingest-jsonl \
--metadata results.json \
base.ingest.jsonl \
results.ingest.jsonl \
"${{ matrix.benchmark.name }}" > comment.md
cat comment.md >> "$GITHUB_STEP_SUMMARY"

- name: Comment PR
Expand Down
33 changes: 26 additions & 7 deletions .github/workflows/sql-benchmarks.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ jobs:
${{ matrix.scale_factor && format('--opt scale-factor={0}', matrix.scale_factor) || '' }}

- name: Capture file sizes
if: matrix.remote_key == null
if: inputs.mode == 'develop' && matrix.remote_key == null
shell: bash
run: |
uv run --no-project scripts/capture-file-sizes.py \
Expand All @@ -201,17 +201,36 @@ jobs:
-o sizes.json
cat sizes.json >> results.json

- name: Configure AWS credentials for RDS baseline (OIDC)
if: inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false
uses: aws-actions/configure-aws-credentials@e6de054238d6b7531b4efff3b6587d9aade6a06c # v6
with:
role-to-assume: ${{ vars.GH_BENCH_INGEST_ROLE_ARN }}
aws-region: ${{ vars.RDS_BENCH_REGION }}

- name: Compare results
if: inputs.mode == 'pr'
if: inputs.mode == 'pr' && github.event.pull_request.head.repo.fork == false
shell: bash
env:
RDS_BENCH_INSTANCE_ENDPOINT: ${{ vars.RDS_BENCH_INSTANCE_ENDPOINT }}
RDS_BENCH_DB_NAME: ${{ vars.RDS_BENCH_DB_NAME }}
AWS_REGION: ${{ vars.RDS_BENCH_REGION }}
run: |
set -Eeu -o pipefail -x

python3 scripts/s3-download.py s3://vortex-ci-benchmark-results/data.json.gz data.json.gz --no-sign-request
gzip -d -c data.json.gz > base.json

uv run --no-project scripts/compare-benchmark-jsons.py base.json results.json "${{ matrix.name }}" \
> comment.md
curl -fsSL https://truststore.pki.rds.amazonaws.com/global/global-bundle.pem \
-o "${RUNNER_TEMP}/rds-global-bundle.pem"
DSN="postgresql://bench_ingest@${RDS_BENCH_INSTANCE_ENDPOINT}:5432/${RDS_BENCH_DB_NAME}?sslmode=verify-full&sslrootcert=${RUNNER_TEMP}/rds-global-bundle.pem"
uv run --no-project scripts/fetch-benchmark-baseline.py results.ingest.jsonl \
--postgres "${DSN}" \
--region "${AWS_REGION}" \
--output base.ingest.jsonl
uv run --no-project scripts/compare-benchmark-jsons.py \
--ingest-jsonl \
--metadata results.json \
base.ingest.jsonl \
results.ingest.jsonl \
"${{ matrix.name }}" > comment.md
cat comment.md >> "$GITHUB_STEP_SUMMARY"

- name: Comment PR
Expand Down
235 changes: 231 additions & 4 deletions scripts/compare-benchmark-jsons.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright the Vortex contributors

import argparse
import math
import os
import re
import sys
from dataclasses import dataclass
from io import StringIO
from pathlib import Path
from typing import Any

import numpy as np
Expand All @@ -43,6 +44,203 @@
Z_SCORE_99 = 2.5758293035489004
CONTROL_FORMAT = "parquet"
FILE_SIZE_METRIC = "file_size"
VORTEX_FORMAT = "vortex-file-compressed"


def read_ingest_records(path: str | Path) -> list[dict[str, Any]]:
"""Read canonical v3 benchmark records from JSONL."""

records = []
with open(path, encoding="utf-8") as lines:
for line_no, line in enumerate(lines, start=1):
if not line.strip():
continue
record = orjson.loads(line)
if not isinstance(record, dict):
raise ValueError(f"{path}:{line_no}: expected a JSON object")
records.append(record)
if not records:
raise ValueError(f"{path}: no benchmark records")
return records


def _dataset_label(record: dict[str, Any]) -> str:
dataset = str(record["dataset"])
variant = record.get("dataset_variant")
return dataset if variant is None else f"{dataset}/{variant}"


def _comparison_row(
record: dict[str, Any],
*,
name: str,
unit: str,
value_field: str,
) -> dict[str, Any]:
row = {
"name": name,
"unit": unit,
"value": record[value_field],
"commit_id": record["commit_sha"],
}
if "all_runtimes_ns" in record:
row["all_runtimes"] = record["all_runtimes_ns"]
return row


def _compression_time_name(record: dict[str, Any]) -> str:
operation = {"encode": "compress", "decode": "decompress"}.get(record["op"], str(record["op"]))
prefix = {
VORTEX_FORMAT: "",
"parquet": "parquet_rs-zstd ",
"lance": "lance ",
}.get(record["format"], f"{record['format']} ")
return f"{prefix}{operation} time/{_dataset_label(record)}"


def _compression_ratio_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Derive the compression ratios intentionally omitted from the v3 facts."""

time_values: dict[tuple[str, str | None, str, str, str], float] = {}
size_values: dict[tuple[str, str | None, str, str], float] = {}
for record in records:
kind = record.get("kind")
if kind == "compression_time":
key = (
str(record["dataset"]),
record.get("dataset_variant"),
str(record["op"]),
str(record["format"]),
str(record["commit_sha"]),
)
time_values[key] = float(record["value_ns"])
elif kind == "compression_size":
key = (
str(record["dataset"]),
record.get("dataset_variant"),
str(record["format"]),
str(record["commit_sha"]),
)
size_values[key] = float(record["value_bytes"])

rows = []
groups = sorted(
{(dataset, variant, commit) for dataset, variant, _format, commit in size_values},
key=lambda group: tuple("" if value is None else str(value) for value in group),
)
for dataset, variant, commit in groups:
label = dataset if variant is None else f"{dataset}/{variant}"
vortex_size = size_values.get((dataset, variant, VORTEX_FORMAT, commit))
for other_format, comparison_name in (("parquet", "parquet-zstd"), ("lance", "lance")):
other_size = size_values.get((dataset, variant, other_format, commit))
if vortex_size is not None and other_size not in (None, 0):
rows.append(
{
"name": f"vortex:{comparison_name} size/{label}",
"unit": "ratio",
"value": vortex_size / other_size,
"commit_id": commit,
}
)

time_groups = sorted(
{(dataset, variant, operation, commit) for dataset, variant, operation, _format, commit in time_values},
key=lambda group: tuple("" if value is None else str(value) for value in group),
)
for dataset, variant, operation, commit in time_groups:
label = dataset if variant is None else f"{dataset}/{variant}"
vortex_time = time_values.get((dataset, variant, operation, VORTEX_FORMAT, commit))
operation_name = {"encode": "compress", "decode": "decompress"}.get(operation, operation)
for other_format, comparison_name in (("parquet", "parquet-zstd"), ("lance", "lance")):
other_time = time_values.get((dataset, variant, operation, other_format, commit))
if vortex_time is not None and other_time not in (None, 0):
rows.append(
{
"name": f"vortex:{comparison_name} ratio {operation_name} time/{label}",
"unit": "ratio",
"value": vortex_time / other_time,
"commit_id": commit,
}
)
return rows


def normalize_ingest_records(records: list[dict[str, Any]]) -> pd.DataFrame:
"""Convert v3 fact records into the comparison reporter's common row shape."""

rows: list[dict[str, Any]] = []
known_kinds = {
"query_measurement",
"compression_time",
"compression_size",
"random_access_time",
"vector_search_run",
}
for record in records:
kind = record.get("kind")
if kind not in known_kinds:
raise ValueError(f"unknown benchmark record kind: {kind!r}")

if kind == "query_measurement":
row = _comparison_row(
record,
name=(f"{record['dataset']}_q{int(record['query_idx']):02}/{record['engine']}:{record['format']}"),
unit="ns",
value_field="value_ns",
)
row["storage"] = record["storage"]
row["dataset"] = {
"dataset": record["dataset"],
"dataset_variant": record.get("dataset_variant"),
"scale_factor": record.get("scale_factor"),
}
rows.append(row)
elif kind == "compression_time":
rows.append(
_comparison_row(
record,
name=_compression_time_name(record),
unit="ns",
value_field="value_ns",
)
)
elif kind == "compression_size":
rows.append(
_comparison_row(
record,
name=f"{record['format']} size/{_dataset_label(record)}",
unit="bytes",
value_field="value_bytes",
)
)
elif kind == "random_access_time":
format_extension = {
VORTEX_FORMAT: "vortex",
"vortex-compact": "vortex",
}.get(record["format"], record["format"])
dataset_path = "" if record["dataset"] == "taxi" else f"{record['dataset']}/"
row = _comparison_row(
record,
name=f"random-access/{dataset_path}{format_extension}-tokio-local-disk",
unit="ns",
value_field="value_ns",
)
row["storage"] = "nvme"
rows.append(row)
else:
rows.append(
_comparison_row(
record,
name=(
f"vector-search/{record['dataset']}/{record['layout']}/{record['flavor']}/{record['threshold']}"
),
unit="ns",
value_field="value_ns",
)
)

rows.extend(_compression_ratio_rows(records))
return pd.DataFrame(rows)


@dataclass
Expand Down Expand Up @@ -839,14 +1037,43 @@ def group_sort_key(group_key: tuple[str, str]) -> tuple[int, int, str, str]:
)


def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Render a benchmark comparison as Markdown.")
parser.add_argument("base", help="Baseline JSONL path.")
parser.add_argument("pr", help="Pull-request JSONL path.")
parser.add_argument("benchmark_name", nargs="?", default="", help="Benchmark display name.")
parser.add_argument(
"--ingest-jsonl",
action="store_true",
help="Read both inputs as canonical v3 ingest records instead of legacy gh-json rows.",
)
parser.add_argument(
"--metadata",
help="Optional legacy gh-json path used only for display metadata such as the documentation link.",
)
return parser.parse_args()


def main() -> None:
"""Render the benchmark comparison markdown used in CI PR comments."""

benchmark_name = sys.argv[3] if len(sys.argv) > 3 else ""
args = parse_args()
benchmark_name = args.benchmark_name

if args.ingest_jsonl:
base = normalize_ingest_records(read_ingest_records(args.base))
pr = normalize_ingest_records(read_ingest_records(args.pr))
if args.metadata is not None:
metadata = pd.read_json(args.metadata, lines=True)
if "doc" in metadata.columns:
docs = metadata["doc"].dropna().unique()
if len(docs) > 0:
pr["doc"] = docs[0]
else:
pr = pd.read_json(args.pr, lines=True)
base = read_latest_baseline_rows(args.base, pr)

pr = pd.read_json(sys.argv[2], lines=True)
title = format_title(benchmark_name, pr)
base = read_latest_baseline_rows(sys.argv[1], pr)

base_commit_id = set(base["commit_id"].unique())
pr_commit_id = set(pr["commit_id"].unique())
Expand Down
Loading
Loading