Run new benchmark script in branch. (#10303)
Summary: Configure CI to run modernised benchmark script Pull Request resolved: https://github.com/facebook/rocksdb/pull/10303 Reviewed By: ramvadiv Differential Revision: D37719116 Pulled By: jay-zhuang fbshipit-source-id: 79ecb1cd0abd4d800c6906ba6673268c2adee10emain
parent
01a2e20299
commit
e637470f64
@ -0,0 +1,165 @@ |
|||||||
|
#!/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 os |
||||||
|
import re |
||||||
|
import shutil |
||||||
|
import subprocess |
||||||
|
import sys |
||||||
|
import logging |
||||||
|
|
||||||
|
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'] |
||||||
|
|
||||||
|
|
||||||
|
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()) |
Loading…
Reference in new issue