From 6db682c7179131f27c7e823086476ac0f49c59e1 Mon Sep 17 00:00:00 2001 From: Gregory Kimball Date: Thu, 23 Jul 2026 04:13:48 +0000 Subject: [PATCH 1/2] Throughput testing fixture --- presto/scripts/run_throughput_benchmark.sh | 269 ++++++++++ .../testing/throughput_benchmarks/__init__.py | 2 + .../testing/throughput_benchmarks/conftest.py | 106 ++++ .../throughput_benchmarks/stream_runner.py | 494 ++++++++++++++++++ .../tpch_throughput_test.py | 140 +++++ 5 files changed, 1011 insertions(+) create mode 100755 presto/scripts/run_throughput_benchmark.sh create mode 100644 presto/testing/throughput_benchmarks/__init__.py create mode 100644 presto/testing/throughput_benchmarks/conftest.py create mode 100644 presto/testing/throughput_benchmarks/stream_runner.py create mode 100644 presto/testing/throughput_benchmarks/tpch_throughput_test.py diff --git a/presto/scripts/run_throughput_benchmark.sh b/presto/scripts/run_throughput_benchmark.sh new file mode 100755 index 00000000..69b7d016 --- /dev/null +++ b/presto/scripts/run_throughput_benchmark.sh @@ -0,0 +1,269 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +set -e + +# Compute the directory where this script resides +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# LOGS_DIR points to the directory where server log files (including nvidia-smi +# output) are written so that run_context.py can parse GPU info. +export LOGS_DIR="${LOGS_DIR:-${SCRIPT_DIR}/presto_logs}" + +source "${SCRIPT_DIR}/presto_connection_defaults.sh" + +print_help() { + cat << EOF + +Usage: $0 [OPTIONS] + +Run a multi-stream TPC throughput benchmark against a running Presto cluster. +Does not drop OS page caches. Verifies ANALYZE TABLE has been run unless +--skip-analyze-check is set. + +OPTIONS: + -h, --help Show this help message. + -b, --benchmark-type Benchmark type. Currently only "tpch" is supported. + -q, --queries Comma-separated allow-list of query numbers (default: all). + --exclude-queries Comma-separated query numbers to exclude (for p99 work). + --queries-file Path to a custom JSON file containing query definitions. + -H, --hostname Hostname of the Presto coordinator. + --port Port number of the Presto coordinator. + -u, --user User who queries will be executed as. + -s, --schema-name Existing schema containing the benchmark tables. + -o, --output-dir Output directory (default: \$(pwd)/throughput_benchmark_output). + -t, --tag Optional tag subdirectory under --output-dir. + --num-streams Number of concurrent client streams / cursors (default: 1). + --suite-repeats Times each stream runs the full (shuffled) suite (default: 1). + --repetitions-per-query Times each query is run before advancing (default: 1). + --run-seed Seed for deterministic per-stream shuffle (default: 1). + --skip-analyze-check Skip ANALYZE TABLE verification. + --reference-results-dir Gold parquet dir (typically Presto C++ query_results/). + --no-save-results Do not write query_results parquet files. + -v, --verbose Print debug logs for worker/engine detection. + +Semantics: + With --num-streams N, --suite-repeats R, --repetitions-per-query P: + each of N streams runs the query suite R times (seeded shuffle per suite), + and within a suite each query is executed P times before moving on. + + Execution retries up to 10 times with no backoff; failed-attempt time is + included in the recorded latency. Validation (if --reference-results-dir is + set) runs in an untimed region after a successful execution. + +EXAMPLES: + $0 -b tpch -s bench_sf100 --num-streams 4 --suite-repeats 5 --run-seed 42 + $0 -b tpch -s bench_sf1000 -q "1,3,6,12" --exclude-queries "15,17" --num-streams 8 + $0 -b tpch -s bench_sf100 --reference-results-dir /path/to/cpu/query_results + +EOF +} + +BENCHMARK_TYPE="" +QUERIES="" +EXCLUDE_QUERIES="" +QUERIES_FILE="" +HOST_NAME="" +PORT="" +USER_NAME="" +SCHEMA_NAME="" +OUTPUT_DIR="" +TAG="" +NUM_STREAMS="" +SUITE_REPEATS="" +REPETITIONS_PER_QUERY="" +RUN_SEED="" +SKIP_ANALYZE_CHECK="" +PRESTO_EXPECTED_RESULTS_DIR="" +EXPLICIT_REFERENCE_DIR="" +NO_SAVE_RESULTS="" + +parse_args() { + while [[ $# -gt 0 ]]; do + case $1 in + -h|--help) + print_help + exit 0 + ;; + -b|--benchmark-type) + BENCHMARK_TYPE=$2; shift 2 + ;; + -q|--queries) + QUERIES=$2; shift 2 + ;; + --exclude-queries) + EXCLUDE_QUERIES=$2; shift 2 + ;; + --queries-file) + QUERIES_FILE=$2; shift 2 + ;; + -H|--hostname) + HOST_NAME=$2; shift 2 + ;; + --port) + PORT=$2; shift 2 + ;; + -u|--user) + USER_NAME=$2; shift 2 + ;; + -s|--schema-name) + SCHEMA_NAME=$2; shift 2 + ;; + -o|--output-dir) + OUTPUT_DIR=$2; shift 2 + ;; + -t|--tag) + TAG=$2; shift 2 + ;; + --num-streams) + NUM_STREAMS=$2; shift 2 + ;; + --suite-repeats) + SUITE_REPEATS=$2; shift 2 + ;; + --repetitions-per-query) + REPETITIONS_PER_QUERY=$2; shift 2 + ;; + --run-seed) + RUN_SEED=$2; shift 2 + ;; + --skip-analyze-check) + SKIP_ANALYZE_CHECK=true; shift + ;; + --reference-results-dir) + PRESTO_EXPECTED_RESULTS_DIR=$2 + EXPLICIT_REFERENCE_DIR=true + shift 2 + ;; + --no-save-results) + NO_SAVE_RESULTS=true; shift + ;; + -v|--verbose) + export PRESTO_BENCHMARK_DEBUG=1 + shift + ;; + *) + echo "Error: Unknown argument $1" + print_help + exit 1 + ;; + esac + done +} + +parse_args "$@" + +if [[ -z ${BENCHMARK_TYPE} || ${BENCHMARK_TYPE} != "tpch" ]]; then + echo "Error: A valid benchmark type is required (currently only tpch). Use -b/--benchmark-type." + print_help + exit 1 +fi + +if [[ -z ${SCHEMA_NAME} ]]; then + echo "Error: A schema name must be set. Use the -s or --schema-name argument." + print_help + exit 1 +fi + +if [[ ${EXPLICIT_REFERENCE_DIR} == true && ! -d ${PRESTO_EXPECTED_RESULTS_DIR} ]]; then + echo "[Validation] Error: --reference-results-dir not found: ${PRESTO_EXPECTED_RESULTS_DIR}" >&2 + exit 1 +fi + +set_presto_coordinator_defaults + +PYTEST_ARGS=("--schema-name ${SCHEMA_NAME}") + +if [[ -n ${QUERIES} ]]; then + PYTEST_ARGS+=("--queries ${QUERIES}") +fi +if [[ -n ${EXCLUDE_QUERIES} ]]; then + PYTEST_ARGS+=("--exclude-queries ${EXCLUDE_QUERIES}") +fi +if [[ -n ${QUERIES_FILE} ]]; then + PYTEST_ARGS+=("--queries-file ${QUERIES_FILE}") +fi +if [[ -n ${HOST_NAME} ]]; then + PYTEST_ARGS+=("--hostname ${HOST_NAME}") +fi +if [[ -n ${PORT} ]]; then + PYTEST_ARGS+=("--port ${PORT}") +fi +if [[ -n ${USER_NAME} ]]; then + PYTEST_ARGS+=("--user ${USER_NAME}") +fi +if [[ -n ${OUTPUT_DIR} ]]; then + PYTEST_ARGS+=("--output-dir ${OUTPUT_DIR}") +fi +if [[ -n ${TAG} ]]; then + if [[ ! ${TAG} =~ ^[a-zA-Z0-9_]+$ ]]; then + echo "Error: Invalid --tag value. Tags must contain only alphanumeric and underscore characters." + exit 1 + fi + PYTEST_ARGS+=("--tag ${TAG}") +fi +if [[ -n ${NUM_STREAMS} ]]; then + PYTEST_ARGS+=("--num-streams ${NUM_STREAMS}") +fi +if [[ -n ${SUITE_REPEATS} ]]; then + PYTEST_ARGS+=("--suite-repeats ${SUITE_REPEATS}") +fi +if [[ -n ${REPETITIONS_PER_QUERY} ]]; then + PYTEST_ARGS+=("--repetitions-per-query ${REPETITIONS_PER_QUERY}") +fi +if [[ -n ${RUN_SEED} ]]; then + PYTEST_ARGS+=("--run-seed ${RUN_SEED}") +fi +if [[ "${SKIP_ANALYZE_CHECK}" == "true" ]]; then + PYTEST_ARGS+=("--skip-analyze-check") +fi +if [[ -n ${PRESTO_EXPECTED_RESULTS_DIR} ]]; then + PYTEST_ARGS+=("--reference-results-dir ${PRESTO_EXPECTED_RESULTS_DIR}") +fi +if [[ "${NO_SAVE_RESULTS}" == "true" ]]; then + PYTEST_ARGS+=("--no-save-results") +fi + +source "${SCRIPT_DIR}/../../scripts/py_env_functions.sh" + +trap delete_python_virtual_env EXIT + +init_python_virtual_env + +TEST_DIR=$(readlink -f "${SCRIPT_DIR}/../testing") +pip install -q -r ${TEST_DIR}/requirements.txt + +source "${SCRIPT_DIR}/common_functions.sh" + +wait_for_worker_node_registration "$HOST_NAME" "$PORT" + +echo "Running throughput bench" +echo " schema=${SCHEMA_NAME} streams=${NUM_STREAMS:-1} suite_repeats=${SUITE_REPEATS:-1} repetitions=${REPETITIONS_PER_QUERY:-1} seed=${RUN_SEED:-1}" + +THROUGHPUT_TEST_DIR=${TEST_DIR}/throughput_benchmarks +# Ensure repo root is importable (common.* / presto.testing.*) +export PYTHONPATH="$(readlink -f "${SCRIPT_DIR}/../.."):${PYTHONPATH:-}" + +pytest -q -s ${THROUGHPUT_TEST_DIR}/tpch_throughput_test.py ${PYTEST_ARGS[*]} + +EFFECTIVE_OUTPUT_DIR="$(readlink -f "${OUTPUT_DIR:-$(pwd)/throughput_benchmark_output}")" +if [[ -n "${TAG}" ]]; then + EFFECTIVE_BENCHMARK_DIR="${EFFECTIVE_OUTPUT_DIR}/${TAG}" +else + EFFECTIVE_BENCHMARK_DIR="${EFFECTIVE_OUTPUT_DIR}" +fi + +if [[ -d "${LOGS_DIR}" ]]; then + mkdir -p "${EFFECTIVE_BENCHMARK_DIR}/logs" + cp -r "${LOGS_DIR}/." "${EFFECTIVE_BENCHMARK_DIR}/logs/" || true + echo "Snapshotted logs to ${EFFECTIVE_BENCHMARK_DIR}/logs" +fi + +CONFIG_SRC="$(readlink -f "${SCRIPT_DIR}/../docker/config/generated")" +if [[ -d "${CONFIG_SRC}" ]]; then + mkdir -p "${EFFECTIVE_BENCHMARK_DIR}/config" + cp -r "${CONFIG_SRC}/." "${EFFECTIVE_BENCHMARK_DIR}/config/" || true + echo "Snapshotted configs to ${EFFECTIVE_BENCHMARK_DIR}/config" +fi diff --git a/presto/testing/throughput_benchmarks/__init__.py b/presto/testing/throughput_benchmarks/__init__.py new file mode 100644 index 00000000..e37bf1e8 --- /dev/null +++ b/presto/testing/throughput_benchmarks/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 diff --git a/presto/testing/throughput_benchmarks/conftest.py b/presto/testing/throughput_benchmarks/conftest.py new file mode 100644 index 00000000..a6ab25f2 --- /dev/null +++ b/presto/testing/throughput_benchmarks/conftest.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from datetime import datetime, timezone + +import prestodb +import pytest + +from common.testing.performance_benchmarks.benchmark_keys import BenchmarkKeys +from common.testing.performance_benchmarks.conftest import DataLocation + +from ..common.conftest import * # noqa: F403 +from ..common.fixtures import ( + tpcds_queries, # noqa: F401 + tpch_queries, # noqa: F401 +) +from ..integration_tests.analyze_tables import check_tables_analyzed +from ..performance_benchmarks.run_context import gather_run_context +from common.testing.performance_benchmarks.common_fixtures import benchmark_queries # noqa: F401 + + +def pytest_addoption(parser): + parser.addoption("--queries") + parser.addoption("--exclude-queries", default="") + parser.addoption("--queries-file") + parser.addoption("--schema-name", required=True) + parser.addoption("--scale-factor") + parser.addoption("--hostname", default="localhost") + parser.addoption("--port", default=8080, type=int) + parser.addoption("--user", default="test_user") + # Power-run compatibility: unused for throughput timing, kept so shared + # reporting helpers that read --iterations do not explode. + parser.addoption("--iterations", default=1, type=int) + parser.addoption("--output-dir", default="throughput_benchmark_output") + parser.addoption("--tag") + parser.addoption("--skip-analyze-check", action="store_true", default=False) + parser.addoption("--num-streams", default=1, type=int) + parser.addoption("--suite-repeats", default=1, type=int) + parser.addoption("--repetitions-per-query", default=1, type=int) + parser.addoption("--run-seed", default=1, type=int) + parser.addoption("--reference-results-dir", default=None) + parser.addoption( + "--no-save-results", + action="store_true", + default=False, + help="Do not write query_results parquet files.", + ) + + +def pytest_configure(config): + pytest.data_location = DataLocation("--schema-name", "Schema", BenchmarkKeys.SCHEMA_NAME_KEY) + + +@pytest.fixture(scope="session", autouse=True) +def run_context_collector(request): + hostname = request.config.getoption("--hostname") + port = request.config.getoption("--port") + user = request.config.getoption("--user") + schema_name = request.config.getoption("--schema-name") + + ctx = gather_run_context( + hostname=hostname, + port=port, + user=user, + schema_name=schema_name, + ) + ctx["timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + ctx["mode"] = "throughput" + ctx["concurrency_streams"] = request.config.getoption("--num-streams") + ctx["suite_repeats"] = request.config.getoption("--suite-repeats") + ctx["repetitions_per_query"] = request.config.getoption("--repetitions-per-query") + ctx["run_seed"] = request.config.getoption("--run-seed") + yield ctx + request.session.run_context = ctx + + +@pytest.fixture(scope="session", autouse=True) +def verify_tables_analyzed(request): + """Verify ANALYZE TABLE has been run before collecting times.""" + if request.config.getoption("--skip-analyze-check"): + print("[Analyze] Skipping analyze check (--skip-analyze-check flag set).") + return + hostname = request.config.getoption("--hostname") + port = request.config.getoption("--port") + user = request.config.getoption("--user") + schema = request.config.getoption("--schema-name") + conn = prestodb.dbapi.connect(host=hostname, port=port, user=user, catalog="hive", schema=schema) + cursor = conn.cursor() + try: + check_tables_analyzed(cursor, schema) + except RuntimeError as e: + pytest.exit(str(e), returncode=1) + finally: + cursor.close() + conn.close() + + +@pytest.fixture(scope="module") +def presto_cursor(request): + """Cursor used only for query setup (Q11 scale factor); streams open their own.""" + hostname = request.config.getoption("--hostname") + port = request.config.getoption("--port") + user = request.config.getoption("--user") + schema = request.config.getoption("--schema-name") + conn = prestodb.dbapi.connect(host=hostname, port=port, user=user, catalog="hive", schema=schema) + return conn.cursor() diff --git a/presto/testing/throughput_benchmarks/stream_runner.py b/presto/testing/throughput_benchmarks/stream_runner.py new file mode 100644 index 00000000..0aefa56e --- /dev/null +++ b/presto/testing/throughput_benchmarks/stream_runner.py @@ -0,0 +1,494 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-stream TPC throughput runner. + +Each stream owns a dedicated Presto connection/cursor and executes the full +query suite sequentially. Suite order is seeded and reshuffled per suite +repeat. Query execution retries until success (no backoff); wall time of +failed attempts is included in the measured latency. Validation runs in an +untimed region after a successful execution. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import statistics +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pandas as pd +import prestodb + +from common.testing.result_comparison import ValidationStatus, validate_query_result + +MAX_EXECUTION_ATTEMPTS = 10 +MAX_RECORDED_ERRORS = 3 +ERROR_SNIPPET_CHARS = 100 + + +def derive_stream_seed(run_seed: int, stream_id: int) -> int: + """Stable 32-bit seed derived from the main run seed and stream id.""" + digest = hashlib.sha256(f"{run_seed}:{stream_id}".encode()).digest() + return int.from_bytes(digest[:8], "big") + + +def _snippet(text: str, n: int = ERROR_SNIPPET_CHARS) -> str: + text = text.replace("\n", "\\n") + if len(text) <= 2 * n: + return text + return f"{text[:n]}...{text[-n:]}" + + +def _format_error(exc: BaseException) -> dict[str, str]: + message = str(exc) + error_type = getattr(exc, "error_type", None) + error_name = getattr(exc, "error_name", None) + if error_type or error_name: + message = f"{error_type or '?'}: {error_name or '?'}: {message}" + return { + "type": type(exc).__name__, + "message": message, + "first_100": message[:ERROR_SNIPPET_CHARS], + "last_100": message[-ERROR_SNIPPET_CHARS:], + "snippet": _snippet(message), + } + + +@dataclass +class AttemptRecord: + attempt: int + elapsed_ms: float + ok: bool + error: dict[str, str] | None = None + + +@dataclass +class QueryExecution: + query_id: str + suite_repeat: int + repetition: int + total_elapsed_ms: float + status: str # ok | execution_fail | validation_fail + attempts: list[AttemptRecord] = field(default_factory=list) + validation_status: ValidationStatus | None = None + validation_message: str | None = None + + +@dataclass +class StreamResult: + stream_id: int + seed: int + executions: list[QueryExecution] = field(default_factory=list) + errors: list[dict[str, str]] = field(default_factory=list) + wall_ms: float = 0.0 + + +class StreamLogger: + """One log file per stream/cursor.""" + + def __init__(self, path: Path, stream_id: int, seed: int): + self.path = path + self.path.parent.mkdir(parents=True, exist_ok=True) + self._logger = logging.getLogger(f"throughput.stream.{stream_id}") + self._logger.setLevel(logging.INFO) + self._logger.handlers.clear() + self._logger.propagate = False + handler = logging.FileHandler(self.path, mode="w") + handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s")) + self._logger.addHandler(handler) + self._logger.info("stream_id=%s seed=%s", stream_id, seed) + + def info(self, msg: str, *args: Any) -> None: + self._logger.info(msg, *args) + + def error(self, msg: str, *args: Any) -> None: + self._logger.error(msg, *args) + + def close(self) -> None: + for handler in list(self._logger.handlers): + handler.close() + self._logger.removeHandler(handler) + + +def _shuffled_query_order(query_ids: list[str], seed: int, suite_repeat: int) -> list[str]: + import random + + order = list(query_ids) + rng = random.Random(seed + suite_repeat * 1_000_003) + rng.shuffle(order) + return order + + +def _execute_with_retries( + cursor: Any, + sql: str, + query_id: str, + benchmark_type: str, + stream_logger: StreamLogger, + recorded_errors: list[dict[str, str]], +) -> tuple[Any | None, list[AttemptRecord], float, str]: + """Run query until success or MAX_EXECUTION_ATTEMPTS. + + Returns (cursor_after_success_or_None, attempts, total_elapsed_ms, status). + total_elapsed_ms includes failed attempts (no backoff sleeps). + """ + attempts: list[AttemptRecord] = [] + total_elapsed_ms = 0.0 + tagged_sql = f"--{benchmark_type}_{query_id}_throughput--\n{sql}" + + for attempt in range(1, MAX_EXECUTION_ATTEMPTS + 1): + t0 = time.perf_counter() + try: + result_cursor = cursor.execute(tagged_sql) + # Prefer server-reported elapsed time when available; fall back to client wall. + server_ms = None + stats = getattr(result_cursor, "stats", None) or getattr(cursor, "stats", None) + if isinstance(stats, dict): + server_ms = stats.get("elapsedTimeMillis") + client_ms = (time.perf_counter() - t0) * 1000.0 + elapsed_ms = float(server_ms) if server_ms is not None else client_ms + total_elapsed_ms += elapsed_ms + attempts.append(AttemptRecord(attempt=attempt, elapsed_ms=elapsed_ms, ok=True)) + stream_logger.info( + "%s attempt=%s ok elapsed_ms=%.2f total_ms=%.2f", + query_id, + attempt, + elapsed_ms, + total_elapsed_ms, + ) + return result_cursor, attempts, total_elapsed_ms, "ok" + except Exception as exc: # noqa: BLE001 — record and retry any client/server failure + client_ms = (time.perf_counter() - t0) * 1000.0 + total_elapsed_ms += client_ms + err = _format_error(exc) + attempts.append(AttemptRecord(attempt=attempt, elapsed_ms=client_ms, ok=False, error=err)) + if len(recorded_errors) < MAX_RECORDED_ERRORS: + recorded_errors.append(err) + stream_logger.error( + "%s attempt=%s fail elapsed_ms=%.2f error=%s", + query_id, + attempt, + client_ms, + err["snippet"], + ) + + return None, attempts, total_elapsed_ms, "execution_fail" + + +def _maybe_validate( + df: pd.DataFrame, + query_id: str, + query_sql: str, + reference_dir: Path | None, +) -> tuple[ValidationStatus, str | None]: + if reference_dir is None: + return "not-validated", "no reference results directory" + q_num = int(query_id.lstrip("Qq")) + expected_file = next( + ( + reference_dir / name + for name in (f"q{q_num:02d}.parquet", f"q{q_num}.parquet", f"{q_num:02d}.parquet") + if (reference_dir / name).exists() + ), + None, + ) + if expected_file is None: + return "not-validated", f"expected parquet not found for {query_id}" + expected = pd.read_parquet(expected_file) + return validate_query_result(query_id.lower(), df, expected, query_sql) + + +def run_stream( + stream_id: int, + run_seed: int, + hostname: str, + port: int, + user: str, + schema: str, + benchmark_type: str, + queries: dict[str, str], + query_ids: list[str], + suite_repeats: int, + repetitions_per_query: int, + output_dir: Path, + reference_results_dir: Path | None, + save_results: bool, + save_lock: threading.Lock, + saved_queries: set[str], +) -> StreamResult: + seed = derive_stream_seed(run_seed, stream_id) + log_path = output_dir / "stream_logs" / f"stream_{stream_id}.log" + stream_logger = StreamLogger(log_path, stream_id, seed) + result = StreamResult(stream_id=stream_id, seed=seed) + results_dir = output_dir / "query_results" + + wall_t0 = time.perf_counter() + conn = prestodb.dbapi.connect(host=hostname, port=port, user=user, catalog="hive", schema=schema) + cursor = conn.cursor() + try: + for suite_repeat in range(suite_repeats): + order = _shuffled_query_order(query_ids, seed, suite_repeat) + stream_logger.info("suite_repeat=%s order=%s", suite_repeat, ",".join(order)) + for query_id in order: + sql = queries[query_id] + for repetition in range(repetitions_per_query): + result_cursor, attempts, total_ms, status = _execute_with_retries( + cursor=cursor, + sql=sql, + query_id=query_id, + benchmark_type=benchmark_type, + stream_logger=stream_logger, + recorded_errors=result.errors, + ) + execution = QueryExecution( + query_id=query_id, + suite_repeat=suite_repeat, + repetition=repetition, + total_elapsed_ms=total_ms, + status=status, + attempts=attempts, + ) + if status != "ok" or result_cursor is None: + result.executions.append(execution) + continue + + # Untimed region: fetch + optional validation / result capture. + rows = result_cursor.fetchall() + columns = [desc[0] for desc in result_cursor.description] + df = pd.DataFrame(rows, columns=columns) + + if save_results: + with save_lock: + if query_id not in saved_queries: + results_dir.mkdir(parents=True, exist_ok=True) + parquet_path = results_dir / f"{query_id.lower()}.parquet" + if not parquet_path.exists(): + df.to_parquet(parquet_path, index=False) + saved_queries.add(query_id) + + val_status, val_message = _maybe_validate( + df, query_id, sql, reference_results_dir + ) + execution.validation_status = val_status + execution.validation_message = val_message + if val_status == "failed": + execution.status = "validation_fail" + stream_logger.error( + "%s validation_fail message=%s", query_id, _snippet(val_message or "") + ) + else: + stream_logger.info( + "%s validation=%s message=%s", + query_id, + val_status, + _snippet(val_message or ""), + ) + + result.executions.append(execution) + finally: + try: + cursor.close() + except Exception: # noqa: BLE001 + pass + try: + conn.close() + except Exception: # noqa: BLE001 + pass + result.wall_ms = (time.perf_counter() - wall_t0) * 1000.0 + stream_logger.info("stream wall_ms=%.2f executions=%s", result.wall_ms, len(result.executions)) + stream_logger.close() + + # Persist structured stream summary beside the text log. + summary_path = output_dir / "stream_logs" / f"stream_{stream_id}.json" + summary_path.write_text( + json.dumps( + { + "stream_id": stream_id, + "seed": seed, + "wall_ms": result.wall_ms, + "errors": result.errors, + "executions": [ + { + "query_id": e.query_id, + "suite_repeat": e.suite_repeat, + "repetition": e.repetition, + "total_elapsed_ms": e.total_elapsed_ms, + "status": e.status, + "validation_status": e.validation_status, + "validation_message": e.validation_message, + "attempts": [ + { + "attempt": a.attempt, + "elapsed_ms": a.elapsed_ms, + "ok": a.ok, + "error": a.error, + } + for a in e.attempts + ], + } + for e in result.executions + ], + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + return result + + +def run_throughput( + *, + num_streams: int, + run_seed: int, + hostname: str, + port: int, + user: str, + schema: str, + benchmark_type: str, + queries: dict[str, str], + query_ids: list[str], + suite_repeats: int, + repetitions_per_query: int, + output_dir: Path, + reference_results_dir: Path | None, + save_results: bool = True, +) -> dict[str, Any]: + """Launch concurrent streams and aggregate throughput metrics.""" + output_dir.mkdir(parents=True, exist_ok=True) + overall_t0 = time.perf_counter() + stream_results: list[StreamResult] = [] + save_lock = threading.Lock() + saved_queries: set[str] = set() + + with ThreadPoolExecutor(max_workers=num_streams) as pool: + futures = { + pool.submit( + run_stream, + stream_id=stream_id, + run_seed=run_seed, + hostname=hostname, + port=port, + user=user, + schema=schema, + benchmark_type=benchmark_type, + queries=queries, + query_ids=query_ids, + suite_repeats=suite_repeats, + repetitions_per_query=repetitions_per_query, + output_dir=output_dir, + reference_results_dir=reference_results_dir, + save_results=save_results, + save_lock=save_lock, + saved_queries=saved_queries, + ): stream_id + for stream_id in range(num_streams) + } + for future in as_completed(futures): + stream_results.append(future.result()) + + overall_wall_ms = (time.perf_counter() - overall_t0) * 1000.0 + stream_results.sort(key=lambda s: s.stream_id) + return aggregate_results( + stream_results=stream_results, + query_ids=query_ids, + overall_wall_ms=overall_wall_ms, + num_streams=num_streams, + suite_repeats=suite_repeats, + repetitions_per_query=repetitions_per_query, + run_seed=run_seed, + ) + + +def _percentile(values: list[float], p: float) -> float | None: + if not values: + return None + if len(values) == 1: + return values[0] + ordered = sorted(values) + # Nearest-rank style for p99 with small samples. + idx = min(len(ordered) - 1, max(0, int(round(p * (len(ordered) - 1))))) + return ordered[idx] + + +def aggregate_results( + *, + stream_results: list[StreamResult], + query_ids: list[str], + overall_wall_ms: float, + num_streams: int, + suite_repeats: int, + repetitions_per_query: int, + run_seed: int, +) -> dict[str, Any]: + raw_times: dict[str, list[float]] = {qid: [] for qid in query_ids} + stream_times: dict[str, dict[str, list[float]]] = {} + failed_queries: dict[str, list[dict[str, Any]]] = {} + validation: dict[str, dict[str, Any]] = {} + + completed = 0 + for stream in stream_results: + sid = str(stream.stream_id) + stream_times[sid] = {qid: [] for qid in query_ids} + for execution in stream.executions: + qid = execution.query_id + if execution.status == "ok": + raw_times.setdefault(qid, []).append(execution.total_elapsed_ms) + stream_times[sid].setdefault(qid, []).append(execution.total_elapsed_ms) + completed += 1 + else: + failed_queries.setdefault(qid, []).append( + { + "stream_id": stream.stream_id, + "suite_repeat": execution.suite_repeat, + "repetition": execution.repetition, + "status": execution.status, + "validation_status": execution.validation_status, + "validation_message": execution.validation_message, + "total_elapsed_ms": execution.total_elapsed_ms, + } + ) + if execution.validation_status and execution.validation_status != "not-validated": + validation[qid] = { + "status": execution.validation_status, + "message": execution.validation_message, + } + + p99 = {qid: _percentile(times, 0.99) for qid, times in raw_times.items()} + p50 = {qid: statistics.median(times) if times else None for qid, times in raw_times.items()} + avg = {qid: statistics.mean(times) if times else None for qid, times in raw_times.items()} + + queries_per_sec = (completed / (overall_wall_ms / 1000.0)) if overall_wall_ms > 0 else 0.0 + + return { + "raw_times_ms": raw_times, + "stream_times_ms": stream_times, + "failed_queries": failed_queries, + "validation": validation, + "agg_times_ms": { + "avg": avg, + "median": p50, + "p99": p99, + }, + "throughput": { + "total_wall_ms": overall_wall_ms, + "queries_completed": completed, + "queries_per_sec": queries_per_sec, + "stream_wall_ms": {str(s.stream_id): s.wall_ms for s in stream_results}, + "stream_seeds": {str(s.stream_id): s.seed for s in stream_results}, + }, + "errors_by_stream": {str(s.stream_id): s.errors for s in stream_results}, + "config": { + "concurrency_streams": num_streams, + "suite_repeats": suite_repeats, + "repetitions_per_query": repetitions_per_query, + "run_seed": run_seed, + }, + } diff --git a/presto/testing/throughput_benchmarks/tpch_throughput_test.py b/presto/testing/throughput_benchmarks/tpch_throughput_test.py new file mode 100644 index 00000000..3ccf6c8b --- /dev/null +++ b/presto/testing/throughput_benchmarks/tpch_throughput_test.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +"""TPC-H multi-stream throughput benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from common.testing.conftest import _format_query_ids, _parse_selected_query_ids +from common.testing.performance_benchmarks.benchmark_keys import BenchmarkKeys +from common.testing.performance_benchmarks.conftest import get_output_dir + +from .stream_runner import run_throughput + + +BENCHMARK_TYPE = "tpch" +TPCH_NUM_QUERIES = 22 + + +def _selected_query_ids(config) -> list[str]: + selected = _parse_selected_query_ids(config.getoption("--queries"), TPCH_NUM_QUERIES) + if not selected: + selected = list(range(1, TPCH_NUM_QUERIES + 1)) + + exclude_raw = config.getoption("--exclude-queries") or "" + excluded = set(_parse_selected_query_ids(exclude_raw, TPCH_NUM_QUERIES)) if exclude_raw.strip() else set() + selected = [qid for qid in selected if qid not in excluded] + if not selected: + raise pytest.UsageError("No queries left after applying --queries / --exclude-queries") + return _format_query_ids(selected) + + +def test_tpch_throughput(request, benchmark_queries, run_context_collector): + config = request.config + query_ids = _selected_query_ids(config) + output_dir = get_output_dir(config) + output_dir.mkdir(parents=True, exist_ok=True) + + reference = config.getoption("--reference-results-dir") + reference_dir = Path(reference) if reference else None + if reference_dir is not None and not reference_dir.is_dir(): + raise pytest.UsageError(f"--reference-results-dir not found: {reference_dir}") + + save_results = not config.getoption("--no-save-results") + + aggregate = run_throughput( + num_streams=config.getoption("--num-streams"), + run_seed=config.getoption("--run-seed"), + hostname=config.getoption("--hostname"), + port=config.getoption("--port"), + user=config.getoption("--user"), + schema=config.getoption("--schema-name"), + benchmark_type=BENCHMARK_TYPE, + queries=benchmark_queries, + query_ids=query_ids, + suite_repeats=config.getoption("--suite-repeats"), + repetitions_per_query=config.getoption("--repetitions-per-query"), + output_dir=output_dir, + reference_results_dir=reference_dir, + save_results=save_results, + ) + + json_result = { + BenchmarkKeys.CONTEXT_KEY: { + BenchmarkKeys.ITERATIONS_COUNT_KEY: config.getoption("--iterations"), + BenchmarkKeys.SCHEMA_NAME_KEY: config.getoption("--schema-name"), + "benchmark": [BENCHMARK_TYPE], + "mode": "throughput", + "concurrency_streams": aggregate["config"]["concurrency_streams"], + "suite_repeats": aggregate["config"]["suite_repeats"], + "repetitions_per_query": aggregate["config"]["repetitions_per_query"], + "run_seed": aggregate["config"]["run_seed"], + "query_ids": query_ids, + }, + BENCHMARK_TYPE: { + BenchmarkKeys.RAW_TIMES_KEY: aggregate["raw_times_ms"], + BenchmarkKeys.FAILED_QUERIES_KEY: aggregate["failed_queries"], + BenchmarkKeys.AGGREGATE_TIMES_KEY: aggregate["agg_times_ms"], + "stream_times_ms": aggregate["stream_times_ms"], + "throughput": aggregate["throughput"], + "validation": aggregate["validation"], + "errors_by_stream": aggregate["errors_by_stream"], + }, + } + + tag = config.getoption("--tag") + if tag: + json_result[BenchmarkKeys.CONTEXT_KEY][BenchmarkKeys.TAG_KEY] = tag + + for key, value in run_context_collector.items(): + json_result[BenchmarkKeys.CONTEXT_KEY][key] = value + + result_path = output_dir / "benchmark_result.json" + result_path.write_text(json.dumps(json_result, indent=2) + "\n", encoding="utf-8") + + # Human-readable summary + tp = aggregate["throughput"] + print("") + print("=" * 72) + print("TPC-H Throughput Summary") + print("=" * 72) + print(f"Streams: {aggregate['config']['concurrency_streams']}") + print(f"Suite repeats: {aggregate['config']['suite_repeats']}") + print(f"Repetitions/query: {aggregate['config']['repetitions_per_query']}") + print(f"Run seed: {aggregate['config']['run_seed']}") + print(f"Queries completed: {tp['queries_completed']}") + print(f"Total wall (ms): {tp['total_wall_ms']:.2f}") + print(f"Queries/sec: {tp['queries_per_sec']:.4f}") + print("") + print(f"{'Query':^8}|{'Avg(ms)':^14}|{'Median(ms)':^14}|{'P99(ms)':^14}|{'N':^8}") + print("-" * 72) + for qid in query_ids: + times = aggregate["raw_times_ms"].get(qid) or [] + avg = aggregate["agg_times_ms"]["avg"].get(qid) + med = aggregate["agg_times_ms"]["median"].get(qid) + p99 = aggregate["agg_times_ms"]["p99"].get(qid) + avg_s = f"{avg:.2f}" if avg is not None else "NULL" + med_s = f"{med:.2f}" if med is not None else "NULL" + p99_s = f"{p99:.2f}" if p99 is not None else "NULL" + print(f"{qid:^8}|{avg_s:^14}|{med_s:^14}|{p99_s:^14}|{len(times):^8}") + + failed = aggregate["failed_queries"] + if failed: + print("") + print(f"Failed query executions: {sum(len(v) for v in failed.values())} across {len(failed)} query ids") + for qid, entries in failed.items(): + print(f" {qid}: {len(entries)} failure(s)") + + print(f"\nWrote {result_path}") + print(f"Per-stream logs under {output_dir / 'stream_logs'}") + + # Soft-fail the pytest if any execution/validation failures occurred. + if failed: + pytest.fail( + f"Throughput run completed with failures for queries: {', '.join(sorted(failed.keys()))}" + ) From 4ac1d4825e472c8ef9486e324b1bb952b558833e Mon Sep 17 00:00:00 2001 From: Gregory Kimball Date: Thu, 23 Jul 2026 04:27:24 +0000 Subject: [PATCH 2/2] refactor out the extra folder --- presto/scripts/run_throughput_benchmark.sh | 14 ++- .../performance_benchmarks/conftest.py | 12 ++ .../stream_runner.py | 0 .../tpch_throughput_test.py | 0 .../testing/throughput_benchmarks/__init__.py | 2 - .../testing/throughput_benchmarks/conftest.py | 106 ------------------ 6 files changed, 23 insertions(+), 111 deletions(-) rename presto/testing/{throughput_benchmarks => performance_benchmarks}/stream_runner.py (100%) rename presto/testing/{throughput_benchmarks => performance_benchmarks}/tpch_throughput_test.py (100%) delete mode 100644 presto/testing/throughput_benchmarks/__init__.py delete mode 100644 presto/testing/throughput_benchmarks/conftest.py diff --git a/presto/scripts/run_throughput_benchmark.sh b/presto/scripts/run_throughput_benchmark.sh index 69b7d016..40e562b9 100755 --- a/presto/scripts/run_throughput_benchmark.sh +++ b/presto/scripts/run_throughput_benchmark.sh @@ -174,7 +174,15 @@ fi set_presto_coordinator_defaults -PYTEST_ARGS=("--schema-name ${SCHEMA_NAME}") +OUTPUT_DIR=${OUTPUT_DIR:-"$(pwd)/throughput_benchmark_output"} + +# Reuse performance_benchmarks/conftest.py while preserving throughput +# semantics: no cache drop and no power-run hot/lukewarm iteration model. +PYTEST_ARGS=( + "--schema-name ${SCHEMA_NAME}" + "--iterations 1" + "--skip-drop-cache" +) if [[ -n ${QUERIES} ]]; then PYTEST_ARGS+=("--queries ${QUERIES}") @@ -242,11 +250,11 @@ wait_for_worker_node_registration "$HOST_NAME" "$PORT" echo "Running throughput bench" echo " schema=${SCHEMA_NAME} streams=${NUM_STREAMS:-1} suite_repeats=${SUITE_REPEATS:-1} repetitions=${REPETITIONS_PER_QUERY:-1} seed=${RUN_SEED:-1}" -THROUGHPUT_TEST_DIR=${TEST_DIR}/throughput_benchmarks +PERFORMANCE_TEST_DIR=${TEST_DIR}/performance_benchmarks # Ensure repo root is importable (common.* / presto.testing.*) export PYTHONPATH="$(readlink -f "${SCRIPT_DIR}/../.."):${PYTHONPATH:-}" -pytest -q -s ${THROUGHPUT_TEST_DIR}/tpch_throughput_test.py ${PYTEST_ARGS[*]} +pytest -q -s ${PERFORMANCE_TEST_DIR}/tpch_throughput_test.py ${PYTEST_ARGS[*]} EFFECTIVE_OUTPUT_DIR="$(readlink -f "${OUTPUT_DIR:-$(pwd)/throughput_benchmark_output}")" if [[ -n "${TAG}" ]]; then diff --git a/presto/testing/performance_benchmarks/conftest.py b/presto/testing/performance_benchmarks/conftest.py index e791501d..cebc0130 100644 --- a/presto/testing/performance_benchmarks/conftest.py +++ b/presto/testing/performance_benchmarks/conftest.py @@ -33,6 +33,7 @@ def pytest_addoption(parser): parser.addoption("--queries") + parser.addoption("--exclude-queries", default="") parser.addoption("--queries-file") # path to a custom JSON file containing query definitions parser.addoption("--schema-name", required=True) parser.addoption("--scale-factor") @@ -47,6 +48,17 @@ def pytest_addoption(parser): parser.addoption("--metrics", action="store_true", default=False) parser.addoption("--skip-drop-cache", action="store_true", default=False) parser.addoption("--skip-analyze-check", action="store_true", default=False) + parser.addoption("--num-streams", default=1, type=int) + parser.addoption("--suite-repeats", default=1, type=int) + parser.addoption("--repetitions-per-query", default=1, type=int) + parser.addoption("--run-seed", default=1, type=int) + parser.addoption("--reference-results-dir", default=None) + parser.addoption( + "--no-save-results", + action="store_true", + default=False, + help="Do not write query_results parquet files.", + ) def pytest_configure(config): diff --git a/presto/testing/throughput_benchmarks/stream_runner.py b/presto/testing/performance_benchmarks/stream_runner.py similarity index 100% rename from presto/testing/throughput_benchmarks/stream_runner.py rename to presto/testing/performance_benchmarks/stream_runner.py diff --git a/presto/testing/throughput_benchmarks/tpch_throughput_test.py b/presto/testing/performance_benchmarks/tpch_throughput_test.py similarity index 100% rename from presto/testing/throughput_benchmarks/tpch_throughput_test.py rename to presto/testing/performance_benchmarks/tpch_throughput_test.py diff --git a/presto/testing/throughput_benchmarks/__init__.py b/presto/testing/throughput_benchmarks/__init__.py deleted file mode 100644 index e37bf1e8..00000000 --- a/presto/testing/throughput_benchmarks/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 diff --git a/presto/testing/throughput_benchmarks/conftest.py b/presto/testing/throughput_benchmarks/conftest.py deleted file mode 100644 index a6ab25f2..00000000 --- a/presto/testing/throughput_benchmarks/conftest.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 - -from datetime import datetime, timezone - -import prestodb -import pytest - -from common.testing.performance_benchmarks.benchmark_keys import BenchmarkKeys -from common.testing.performance_benchmarks.conftest import DataLocation - -from ..common.conftest import * # noqa: F403 -from ..common.fixtures import ( - tpcds_queries, # noqa: F401 - tpch_queries, # noqa: F401 -) -from ..integration_tests.analyze_tables import check_tables_analyzed -from ..performance_benchmarks.run_context import gather_run_context -from common.testing.performance_benchmarks.common_fixtures import benchmark_queries # noqa: F401 - - -def pytest_addoption(parser): - parser.addoption("--queries") - parser.addoption("--exclude-queries", default="") - parser.addoption("--queries-file") - parser.addoption("--schema-name", required=True) - parser.addoption("--scale-factor") - parser.addoption("--hostname", default="localhost") - parser.addoption("--port", default=8080, type=int) - parser.addoption("--user", default="test_user") - # Power-run compatibility: unused for throughput timing, kept so shared - # reporting helpers that read --iterations do not explode. - parser.addoption("--iterations", default=1, type=int) - parser.addoption("--output-dir", default="throughput_benchmark_output") - parser.addoption("--tag") - parser.addoption("--skip-analyze-check", action="store_true", default=False) - parser.addoption("--num-streams", default=1, type=int) - parser.addoption("--suite-repeats", default=1, type=int) - parser.addoption("--repetitions-per-query", default=1, type=int) - parser.addoption("--run-seed", default=1, type=int) - parser.addoption("--reference-results-dir", default=None) - parser.addoption( - "--no-save-results", - action="store_true", - default=False, - help="Do not write query_results parquet files.", - ) - - -def pytest_configure(config): - pytest.data_location = DataLocation("--schema-name", "Schema", BenchmarkKeys.SCHEMA_NAME_KEY) - - -@pytest.fixture(scope="session", autouse=True) -def run_context_collector(request): - hostname = request.config.getoption("--hostname") - port = request.config.getoption("--port") - user = request.config.getoption("--user") - schema_name = request.config.getoption("--schema-name") - - ctx = gather_run_context( - hostname=hostname, - port=port, - user=user, - schema_name=schema_name, - ) - ctx["timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - ctx["mode"] = "throughput" - ctx["concurrency_streams"] = request.config.getoption("--num-streams") - ctx["suite_repeats"] = request.config.getoption("--suite-repeats") - ctx["repetitions_per_query"] = request.config.getoption("--repetitions-per-query") - ctx["run_seed"] = request.config.getoption("--run-seed") - yield ctx - request.session.run_context = ctx - - -@pytest.fixture(scope="session", autouse=True) -def verify_tables_analyzed(request): - """Verify ANALYZE TABLE has been run before collecting times.""" - if request.config.getoption("--skip-analyze-check"): - print("[Analyze] Skipping analyze check (--skip-analyze-check flag set).") - return - hostname = request.config.getoption("--hostname") - port = request.config.getoption("--port") - user = request.config.getoption("--user") - schema = request.config.getoption("--schema-name") - conn = prestodb.dbapi.connect(host=hostname, port=port, user=user, catalog="hive", schema=schema) - cursor = conn.cursor() - try: - check_tables_analyzed(cursor, schema) - except RuntimeError as e: - pytest.exit(str(e), returncode=1) - finally: - cursor.close() - conn.close() - - -@pytest.fixture(scope="module") -def presto_cursor(request): - """Cursor used only for query setup (Q11 scale factor); streams open their own.""" - hostname = request.config.getoption("--hostname") - port = request.config.getoption("--port") - user = request.config.getoption("--user") - schema = request.config.getoption("--schema-name") - conn = prestodb.dbapi.connect(host=hostname, port=port, user=user, catalog="hive", schema=schema) - return conn.cursor()