You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
rocksdb/tools/benchmark_ci.py

166 lines
5.4 KiB

#!/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())