wangzhigang1999 commented on code in PR #95: URL: https://github.com/apache/paimon-vector-index/pull/95#discussion_r3951330464
########## tools/benchmark_pr.py: ########## @@ -0,0 +1,382 @@ +#!/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 alert_level(entry): + 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 Review Comment: The report only applies alerts to Recall and QPS; build time, RSS and size have no advisory thresholds. In 2a76386 the helper is renamed query_alert_level and explicitly rejects unsupported directions, with test coverage, so lower-is-better metrics cannot silently receive the wrong interpretation. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
