From e6918187568dbd01842d8d1d2c808ce16a894239 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Sun, 21 Apr 2024 13:54:28 +0200 Subject: Adding upstream version 18.2.2. Signed-off-by: Daniel Baumann --- src/rocksdb/tools/benchmark_ci.py | 182 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 182 insertions(+) create mode 100755 src/rocksdb/tools/benchmark_ci.py (limited to 'src/rocksdb/tools/benchmark_ci.py') diff --git a/src/rocksdb/tools/benchmark_ci.py b/src/rocksdb/tools/benchmark_ci.py new file mode 100755 index 000000000..de9f69cf9 --- /dev/null +++ b/src/rocksdb/tools/benchmark_ci.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# Copyright (c) 2011-present, Facebook, Inc. All rights reserved. +# This source code is licensed under both the GPLv2 (found in the +# COPYING file in the root directory) and Apache 2.0 License +# (found in the LICENSE.Apache file in the root directory). + +"""Run benchmark_compare.sh on the most recent build, for CI +""" + +import argparse +import glob +import logging +import os +import re +import shutil +import subprocess +import sys + +logging.basicConfig(level=logging.INFO) + + +class Config: + def __init__(self, args): + self.version_file = "./include/rocksdb/version.h" + self.data_dir = os.path.expanduser(f"{args.db_dir}") + self.results_dir = os.path.expanduser(f"{args.output_dir}") + self.benchmark_script = f"{os.getcwd()}/tools/benchmark_compare.sh" + self.benchmark_cwd = f"{os.getcwd()}/tools" + + benchmark_env_keys = [ + "LD_LIBRARY_PATH", + "NUM_KEYS", + "KEY_SIZE", + "VALUE_SIZE", + "CACHE_SIZE_MB", + "DURATION_RW", + "DURATION_RO", + "MB_WRITE_PER_SEC", + "NUM_THREADS", + "COMPRESSION_TYPE", + "MIN_LEVEL_TO_COMPRESS", + "WRITE_BUFFER_SIZE_MB", + "TARGET_FILE_SIZE_BASE_MB", + "MAX_BYTES_FOR_LEVEL_BASE_MB", + "MAX_BACKGROUND_JOBS", + "CACHE_INDEX_AND_FILTER_BLOCKS", + "USE_O_DIRECT", + "STATS_INTERVAL_SECONDS", + "SUBCOMPACTIONS", + "COMPACTION_STYLE", + "CI_TESTS_ONLY", + ] + + +def read_version(config): + majorRegex = re.compile(r"#define ROCKSDB_MAJOR\s([0-9]+)") + minorRegex = re.compile(r"#define ROCKSDB_MINOR\s([0-9]+)") + patchRegex = re.compile(r"#define ROCKSDB_PATCH\s([0-9]+)") + with open(config.version_file, "r") as reader: + major = None + minor = None + patch = None + for line in reader: + if major is None: + major = majorRegex.match(line) + elif minor is None: + minor = minorRegex.match(line) + elif patch is None: + patch = patchRegex.match(line) + + if patch is not None: + break + + if patch is not None: + return (major.group(1), minor.group(1), patch.group(1)) + + # Didn't complete a match + return None + + +def prepare(version_str, config): + old_files = glob.glob(f"{config.results_dir}/{version_str}/**", recursive=True) + for f in old_files: + if os.path.isfile(f): + logging.debug(f"remove file {f}") + os.remove(f) + for f in old_files: + if os.path.isdir(f): + logging.debug(f"remove dir {f}") + os.rmdir(f) + + db_bench_vers = f"{config.benchmark_cwd}/db_bench.{version_str}" + + # Create a symlink to the db_bench executable + os.symlink(f"{os.getcwd()}/db_bench", db_bench_vers) + + +def results(version_str, config): + # Copy the report TSV file back to the top level of results + shutil.copyfile( + f"{config.results_dir}/{version_str}/report.tsv", + f"{config.results_dir}/report.tsv", + ) + + +def cleanup(version_str, config): + # Remove the symlink to the db_bench executable + db_bench_vers = f"{config.benchmark_cwd}/db_bench.{version_str}" + os.remove(db_bench_vers) + + +def get_benchmark_env(): + env = [] + for key in Config.benchmark_env_keys: + value = os.getenv(key) + if value is not None: + env.append((key, value)) + return env + + +def main(): + """Tool for running benchmark_compare.sh on the most recent build, for CI + This tool will + + (1) Work out the current version of RocksDB + (2) Run benchmark_compare with that version alone + """ + + parser = argparse.ArgumentParser( + description="benchmark_compare.sh Python wrapper for CI." + ) + + # --tsvfile is the name of the file to read results from + # --esdocument is the ElasticSearch document to push these results into + # + parser.add_argument( + "--db_dir", + default="~/tmp/rocksdb-benchmark-datadir", + help="Database directory hierarchy to use", + ) + parser.add_argument( + "--output_dir", + default="~/tmp/benchmark-results", + help="Benchmark output goes here", + ) + parser.add_argument( + "--num_keys", + default="10000", + help="Number of database keys to use in benchmark test(s) (determines size of test job)", + ) + args = parser.parse_args() + config = Config(args) + + version = read_version(config) + if version is None: + raise Exception(f"Could not read RocksDB version from {config.version_file}") + version_str = f"{version[0]}.{version[1]}.{version[2]}" + logging.info(f"Run benchmark_ci with RocksDB version {version_str}") + + prepare(version_str, config) + + try: + env = get_benchmark_env() + env.append(("NUM_KEYS", args.num_keys)) + cmd = [ + config.benchmark_script, + config.data_dir, + config.results_dir, + version_str, + ] + logging.info(f"Run {cmd} env={env} cwd={config.benchmark_cwd}") + subprocess.run(cmd, env=dict(env), cwd=config.benchmark_cwd) + + results(version_str, config) + finally: + cleanup(version_str, config) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) -- cgit v1.2.3