Optunaを使ったRocksDBのパフォーマンスチューニング

README.md

概要

• Optunaというハイパーパラメータ最適化ツールを使って、RocksDB(組み込みDB・KVS)のパフォーマンスチューニングを試してみた際の結果メモ
  • 対象となるワークロードに対して、最適な性能を発揮するパラメータ群を自動で見つけ出すのが目的
• 結果としては、デフォルトパラメータをそのまま使った場合に比べて、かなり良い性能が得られるパラメータ群を見つけることができた:
  • デフォルトでのベンチマークの所要時間: 372秒
  • Optunaによる最適化後のパラメータでの所要時間: 30秒

モチベーション

• RocksDBには、カスタマイズできるパラメータ群が多数ある（数十～数百？）
  • 自分の利用用途に最適なパラメータ群を人手で調べるのは結構大変
  • RocksDBは、チューニングガイドを含めて、かなりドキュメントが充実しているが、それでも不慣れな人には敷居が高い
  • ↑のような状況は、RocksDBに限らず、ある程度以上規模が大きく歴史が長いミドルウェアだと珍しくない（おそらく）
  • ※ 自分自身が開発者の一人だったり、長年携わっていて習熟度が高い場合は別
• こういった場合、通常は、以下のような流れとなることが多い印象がある（自分が不慣れなミドルウェアの場合）:
  • (1) まずはデフォルト設定で使ってみる
  • (2) それで性能に問題があったら、公式ドキュメントのチューニングガイドを読むなり、Googleで検索するなりして、有名どころのパラメータを調整してみる
  • (3) それでも性能が要求に満たないなら、公式ドキュメントその他を精読し、条件をいろいろと変えつつベンチマークを試行錯誤する
  • (4) デプロイ先の環境が変わったら、それに合わせてまた再調整・再試行
  • => 結構手間なことに加えて、最終的に落ち着いたパラメータ群が最適答かどうかも不明（幅広い組み合わせを試すのはたいていは時間的に無理）
• この辺りのプロセスを極力自動化して、手軽に最適な（ものに近い）パラメータ群を得られるようにしたい、というのがモチベーション
  • とはいえこれはあくまでも理想であって今回は、Optunaがこういった用途に適用可能かどうかを軽く試してみただけ

使用したツール群

• ハイパーパラメータ最適化ツール: optuna-v0.6.0
• ベンチマークツール: ekvsb-v0.0.12
  • 内部的にはRocksDBのRustバインディングのバージョンv0.11.0を使用している

最適化対象のワークロード

今回は、以下のようなかなり単純なワークロードを採用した:

• (1) 50万件分のエントリを保存(PUT)する
  • 値のサイズは10KBに固定 (KVSとしては若干大きめのサイズ）
  • 合計サイズ(論理)は5GB程度
• (2) 上記50万件分のエントリを順不同で取得(GET)する
• (3) 上記50万件分のエントリを順不同で削除(DELETE)する

上記の処理を行うのに適したRocksDBのパラメータ群をOptunaを使って見つけるのが目的。 なお、保存先デバイスにはHDDを使用（RocksDBはデフォルトではSSD向きにチューニングされているので、HDDとは相性が良くない）。

チューニング対象のパラメータ群

RocksDBのRustバインディングが提供しているオプション群は基本的に全てチューニング対象にしている（明らかに趣旨から外れるものは例外; e.g., enable_statistics, set_wal_dir)。

具体的には、以下の通り:

$ ekvsb run rocksdb -h
ekvsb-run-rocksdb 0.0.12

USAGE:
    ekvsb run rocksdb [FLAGS] [OPTIONS] <dir>

FLAGS:
        --block-opt-bloom-filter-block-based
        --block-opt-cache-index-and-filter-blocks
        --block-opt-disable-cache
        --disable-advise-random-on-open
        --disable-auto-compactions
        --disable-concurrent-memtable-write
        --memtable-factory-vector
        --use-direct-io-for-flush-and-compaction
        --use-direct-reads
        --use-fsync

OPTIONS:
        --block-opt-block-size <block-opt-block-size>
        --block-opt-bloom-filter-bits-per-key <block-opt-bloom-filter-bits-per-key>
        --block-opt-index-type <block-opt-index-type>
        --block-opt-lru-cache <block-opt-lru-cache>
        --bytes-per-sync <bytes-per-sync>
        --compaction-readahead-size <compaction-readahead-size>
        --compaction-style <compaction-style> [possible values: Level, Universal, Fifo]
        --level-zero-file-num-compaction-trigger <level-zero-file-num-compaction-trigger>
        --level-zero-slowdown-writes-trigger <level-zero-slowdown-writes-trigger>
        --level-zero-stop-writes-trigger <level-zero-stop-writes-trigger>
        --max-background-compactions <max-background-compactions>
        --max-background-flushes <max-background-flushes>
        --max-bytes-for-level-base <max-bytes-for-level-base>
        --max-bytes-for-level-multiplier <max-bytes-for-level-multiplier>
        --max-manifest-file-size <max-manifest-file-size>
        --max-write-buffer-number <max-write-buffer-number>
        --memtable-factory-hashlinklist-bucket-count <memtable-factory-hashlinklist-bucket-count>
        --memtable-factory-hashskiplist-branching-factor <memtable-factory-hashskiplist-branching-factor>
        --memtable-factory-hashskiplist-bucket-count <memtable-factory-hashskiplist-bucket-count>
        --memtable-factory-hashskiplist-height <memtable-factory-hashskiplist-height>
        --memtable-prefix-bloom-ratio <memtable-prefix-bloom-ratio>
        --min-write-buffer-number <min-write-buffer-number>
        --min-write-buffer-number-to-merge <min-write-buffer-number-to-merge>
        --num-levels <num-levels>
        --optimize-for-point-lookup <optimize-for-point-lookup>
        --optimize-level-style-compaction <optimize-level-style-compaction>
        --parallelism <parallelism>
        --table-cache-num-shard-bits <table-cache-num-shard-bits>
        --target-file-size-base <target-file-size-base>
        --write-buffer-size <write-buffer-size>

Rustバインディングが本家の全てのパラメータを公開している訳ではないので、RocksDBを直接使う場合に比べるとかなり数は少なくなっているが、それでも40個程度のパラメータが存在している。

Optunaの設定

詳細は実際のスクリプト(optimize-rocksdb.py)を見て貰った方が正確、かつ、そこまで変わった使い方はしていないので、ここでは特記事項だけを記載:

• SuccessiveHalvingPrunerによる枝刈り:
  • Optunaの各トライアルで、毎回、500,000件分のPUT/GET/DELETEをフルセットで実行していては時間が掛かってしまうので、見込みがないものに関しては早めに枝刈りするようにしている
  • 具体的には、一回のトライアルで「5,000件分、50,000件分、500,000件分」の三パターンに対して、件数が少ない順にベンチマークを実行し、それぞれの終了時点で枝刈り判定が行われる
    • トライアル全体のおよそ九割は最初の5,000件分の段階で捨てられることになる
  • また、デフォルトパラメータの場合に比べて、極端に実行時間が長い場合にも途中で枝刈り（タイムアウト）が発生するにした
• サンプラーにはRandomSamplerを使用
  • デフォルトのTPESamplerだと、今回のように枝刈りが多数発生する場合には相性が悪いことがあるため (optuna#270)
• 最適化は四時間実行する

実行コマンド

# ワークロードファイルを生成
$ ./gen-workload.sh
$ ls workload.*.json
workload.500000.json  workload.50000.json  workload.5000.json

# パラメータ最適化＆結果
$ python optimize-rocksdb.py
Study statistics:
  Number of finished trials:  938
  Number of pruned trials:  848
  Number of timeout trials: 88
  Number of complete trials:  2
Best trial:
  Value: 29.064814499810876
  Params:
    block-opt-cache-index-and-filter-blocks: 1
    max-write-buffer-number: 32
    memtable-prefix-bloom-ratio: 0.921874416314259
    bytes-per-sync: 21803784
    table-cache-num-shard-bits: 11
    use-direct-reads: 1
    block-opt-bloom-filter-bits-per-key: 58
    block-opt-index-type: HashSearch
    compaction-style: Fifo
    level-zero-stop-writes-trigger: 40
    level-zero-file-num-compaction-trigger: 11
    use-direct-io-for-flush-and-compaction: 0
    memtable-factory: hashlinklist
    block-opt-block-size: 30460927
    memtable-factory-hashlinklist-bucket-count: 1836285
    optimize-for-point-lookup: 1783
    compaction-readahead-size: 14840331
    target-file-size-base: 197272531
    parallelism: 4
    max-background-flushes: 2
    disable-advise-random-on-open: 0
    min-write-buffer-number-to-merge: 24
    max-bytes-for-level-base: 671227089
    use-fsync: 1
    max-bytes-for-level-multiplier: 24.112568220334577
    block-opt-disable-cache: 0
    min-write-buffer-number: 5
    level-zero-slowdown-writes-trigger: 54
    block-opt-bloom-filter-block-based: 1
    block-opt-lru-cache: 30779059
    max-background-compactions: 8
    num-levels: 24
    disable-auto-compactions: 1
    write-buffer-size: 253269056

総トライアル数	完走数	タイムアウト数	枝刈り数
938	2	88	848

上記は、トライアルの枝刈り数等を表にまとめたものだが、これをみるとトライアルのほとんどが枝刈りされていることが分かる。

なお、枝刈りを無効にして最適化を行った場合には、同じ時間で実行できたトライアル数は39であった（その内の35個は途中でタイムアウト。なおタイムアウトも無くした場合には時間内に完了したトライアル数は2個までに減った）。

結果

デフォルトと最適化後のパラメータでのベンチマーク結果の比較:

# デフォルトパラメータでベンチマークを実行
$ cat workload.500000.json | ekvsb run rocksdb /tmp/rocksdb-default/ | ekvsb summary
{
  "oks": 1500000,
  "errors": 0,
  "existence": {
    "exists": 500000,
    "absents": 0,
    "unknowns": 1000000
  },
  "elapsed": 373.72465610034095, # 所要時間(秒)
  "ops": 4013.65009109077,       # 秒間処理操作数
  "latency": {                   # 各操作のレイテンシ(秒)
    "min": 2.9e-6,
    "median": 0.0000324,
    "p95": 0.0004482,
    "p99": 0.0016204,
    "max": 3.8960246
  }
}

# Optunaで最適化されたパラメータでベンチマークを実行
$ cat workload.500000.json | ekvsb run rocksdb /tmp/rocksdb-default/ --block-opt-cache-index-and-filter-blocks --max-write-buffer-number=32 --memtable-prefix-bloom-ratio=0.921874416314259 --bytes-per-sync=21803784 --table-cache-num-shard-bits=11 --use-direct-reads --block-opt-bloom-filter-bits-per-key=58 --block-opt-index-type=HashSearch --compaction-style=Fifo --level-zero-stop-writes-trigger=40 --level-zero-file-num-compaction-trigger=11 --block-opt-block-size=30460927 --memtable-factory-hashlinklist-bucket-count=1836285 --optimize-for-point-lookup=1783 --compaction-readahead-size=14840331 --target-file-size-base=197272531 --parallelism=4 --max-background-flushes=2 --min-write-buffer-number-to-merge=24 --max-bytes-for-level-base=671227089 --use-fsync --max-bytes-for-level-multiplier=24.112568220334577 --min-write-buffer-number=5 --level-zero-slowdown-writes-trigger=54 --block-opt-bloom-filter-block-based --block-opt-lru-cache=30779059 --max-background-compactions=8 --num-levels=24 --disable-auto-compactions --write-buffer-size=253269056 | ekvsb summary
{
  "oks": 1500000,
  "errors": 0,
  "existence": {
    "exists": 500000,
    "absents": 0,
    "unknowns": 1000000
  },
  "elapsed": 29.651524699936886,
  "ops": 50587.617843584034,
  "latency": {
    "min": 2e-6,
    "median": 0.0000112,
    "p95": 0.0000424,
    "p99": 0.000055,
    "max": 0.4047035
  }
}

Optunaによる最適化によって、所要時間が1/10程度になっていることが分かる。

以下の図は、各操作の所要時間(レイテンシ)をプロットしたもの(左=デフォルト、右=Optuna):

感想

• 枝刈り便利
  • パラメータの組み合わせによっては、平気で一桁二桁実行性能に差が出ることがあったけど、遅い場合でも枝刈りで早めに捨てられたので、最適化全体の実行時間に与える影響が少なくできて良かった
• Optunaに「どのパラメータの影響が大きいか(重要か)」を出してくれる機能があると便利
  • それがあれば、とりあえず最初は指定可能なパラメータを全て指定しておいて、後で徐々に絞り込んでいく、といったことが出来そう
  • 後は「影響度が高いものだけを人手でちゃんと調査する」とか「最終的に採用するパラメータ値から影響度が低いものは除外する」といったこともしたい
• 適切なワークロードを設定可能かどうかが鍵なイメージ
  • 誤った対象に最適化しても仕方がないので
  • ワークロードを上手く設定できて、ベンチマークを機械的に実行する仕組みがあるなら、結構使えそう
• 今回は使わなかったが、studyでDBをストレージのバックエンドとして使用する癖をつけた方が良さそう
  • 後で結果を確認したり再開したり、といったことが手軽にできるので
    • 特に最初の試行錯誤段階だと最適化(study.optimize(...))を途中で中断したくなることも良くある
  • ただし、オンメモリのストレージに比べて、SQLiteのストレージは三倍程度遅かった（手元の環境=WSLが悪い可能性もある）
• 実際の用途では、対象ミドルウェア以外(e.g., OS、ベンチマーククライアント)のパラメータ等をチューニングした方が良い場合もあるので要注意
• 最適化対象のパラメータ数が多いと、良い解を得られるまでは結構長い
  • 今回は四時間回したけれど、もっとパラメータを多くしたら、これでも足りなさそうな印象
• 今回のワークロードはかなり簡単かつ人為的なものなので、より実践的なものにした場合にどのような結果になるのかは気になる
• 今回のようなwallclock-timeを指標にするパフォーマンスチューニングの場合には、対象ミドルウェアに指定してパラメータとは関係のないところで、性能が変化することもあるのが面倒
  • 例えば、今回は最初はGCP上の仮想マシンを使って測定を行っていたら、時間経過に伴い明らかにディスクI/Oのスループットが(1/10以下に)下がっていったり、バラつきが大きくなったりして、複数トライアル間で得られた結果をまともに比較できる状態ではなかった（そのため最終的にはローカルマシンを使うようにした）
  • ↑ほど極端ではないとしても、使用しているOSやファイルシステムの特性によって、トライアルの実行タイミング依存で性能が変わる可能性があることは念頭においておく必要がある（理想を言えば、チューニング対象パラメータ以外の要因による性能変化は完全に排除すべきだけど、現実的にはなかなか難しい）

bench.sh

#!/bin/bash

INPUT_FILE=$1
TIMEOUT=$2
DIR=$3

rm -rf $DIR
mkdir -p $DIR

shift 2

cat $INPUT_FILE | timeout $TIMEOUT ekvsb run rocksdb $@ | ekvsb summary

gen-workload.sh

#! /bin/bash

set -eux

for n in 5000 50000 500000
do
    ekvsb workload put --count $n --value-size 10KiB --seed 0 > put.$n.json
    ekvsb workload get --count $n --seed 0 --shuffle 1 > get.$n.json
    ekvsb workload delete --count $n --seed 0 --shuffle 2 > delete.$n.json
    jq -s '.[0] + .[1] + .[2]' put.$n.json get.$n.json delete.$n.json > workload.$n.json
done

optimize-rocksdb.py

import json
from itertools import chain
import optuna
from optuna.pruners import SuccessiveHalvingPruner
from optuna.samplers import RandomSampler
from optuna.structs import TrialPruned
import subprocess

DIR="/tmp/rocksdb/"
WORKLOAD_SIZED=[5000, 50000, 500000]
WORKLOAD_TIMEOUTS=[4, 40, 400]
STUDY_TIMEOUT=4 * 60 * 60

def suggest_bool(trial, name):
    if bool(trial.suggest_categorical(name, [0, 1])):
        return ['--{}'.format(name)]
    else:
        return []

def suggest_int(trial, name, low, high):
    return ['--{}={}'.format(name, trial.suggest_int(name, low, high))]

def suggest_uniform(trial, name, low, high):
    return ['--{}={}'.format(name, trial.suggest_uniform(name, low, high))]

def suggest_categorical(trial, name, list):
    return ['--{}={}'.format(name, trial.suggest_categorical(name, list))]

def objective(trial):
    # このトライアルで試すRocksDBのパラメータセットを選択
    args = [
        suggest_bool(trial, 'disable-advise-random-on-open'),
        suggest_bool(trial, 'disable-auto-compactions'),
        suggest_bool(trial, 'use-direct-io-for-flush-and-compaction'),
        suggest_bool(trial, 'use-direct-reads'),
        suggest_bool(trial, 'use-fsync'),
        suggest_int(trial, 'bytes-per-sync', 0, 32 * 1024 * 1024),
        suggest_int(trial, 'compaction-readahead-size', 0, 32 * 1024 * 1024),
        suggest_categorical(trial, 'compaction-style', ['Level', 'Universal', 'Fifo']),
        suggest_int(trial, 'level-zero-file-num-compaction-trigger', -1, 16),
        suggest_int(trial, 'level-zero-slowdown-writes-trigger', -1, 64),
        suggest_int(trial, 'level-zero-stop-writes-trigger', 0, 64),
        suggest_int(trial, 'max-background-compactions', 1, 8),
        suggest_int(trial, 'max-background-flushes', 1, 8),
        suggest_int(trial, 'max-bytes-for-level-base', 128 * 1024 * 1024, 1024 * 1024 * 1024),
        suggest_uniform(trial, 'max-bytes-for-level-multiplier', 2.0, 32.0),
        suggest_int(trial, 'max-write-buffer-number', -1, 32),
        suggest_uniform(trial, 'memtable-prefix-bloom-ratio', 0.0, 1.0),
        suggest_int(trial, 'min-write-buffer-number', 1, 32),
        suggest_int(trial, 'min-write-buffer-number-to-merge', 0, 32),
        suggest_int(trial, 'num-levels', 1, 64),
        suggest_int(trial, 'optimize-for-point-lookup', 256, 4096),
        suggest_int(trial, 'parallelism', 1, 4),
        suggest_int(trial, 'table-cache-num-shard-bits', 2, 16),
        suggest_int(trial, 'target-file-size-base', 8 * 1024 * 1024, 256 * 1024 * 1024),
        suggest_int(trial, 'write-buffer-size', 8 * 1024 * 1024, 256 * 1024 * 1024),
        suggest_int(trial, 'block-opt-block-size', 512, 32 * 1024 * 1024),
        suggest_int(trial, 'block-opt-lru-cache', 1 * 1024 * 1024, 32 * 1024 * 1024),
        suggest_int(trial, 'block-opt-bloom-filter-bits-per-key', 0, 64),
        suggest_bool(trial, 'block-opt-bloom-filter-block-based'),
        suggest_categorical(trial, 'block-opt-index-type', ['BinarySearch', 'HashSearch', 'TwoLevelIndexSearch']),
    ]
    disable_block_cache = bool(trial.suggest_int('block-opt-disable-cache', 0, 1))
    if disable_block_cache:
        args.extend([['--block-opt-disable-cache']])
    else:
        args.extend([suggest_bool(trial, 'block-opt-cache-index-and-filter-blocks')])

    memtable_factory = trial.suggest_categorical('memtable-factory',
                                                 ['vector', 'skiplist', 'hashskiplist', 'hashlinklist'])
    if memtable_factory == 'skiplist':
        args.extend([
            suggest_bool(trial, 'disable-concurrent-memtable-write'),
        ])
    elif memtable_factory == 'hashskiplist':
        args.extend([
            suggest_int(trial, 'memtable-factory-hashskiplist-bucket-count', 64 * 1024, 2 * 1024 * 1024),
            suggest_int(trial, 'memtable-factory-hashskiplist-height', 2, 16),
            suggest_int(trial, 'memtable-factory-hashskiplist-branching-factor', 2, 16)
        ])
    elif memtable_factory == 'hashlinklist':
        args.extend([
            suggest_int(trial, 'memtable-factory-hashlinklist-bucket-count', 64 * 1024, 2 * 1024 * 1024)
        ])

    # ベンチマークを実行
    for (size, timeout) in zip(WORKLOAD_SIZED, WORKLOAD_TIMEOUTS):
        input_file = "workload.{}.json".format(size)
        proc = subprocess.run(
            ["./bench.sh", input_file, str(timeout), DIR] + list(chain.from_iterable(args)),
            stdout = subprocess.PIPE, check=True)
        result = json.loads(proc.stdout.decode('utf-8'))
        elapsed = result['elapsed']

        trial.report(elapsed, size)
        if trial.should_prune(size):
            raise TrialPruned("size={}".format(size))

    return elapsed

if __name__ == '__main__':
    study = optuna.create_study(
        pruner=SuccessiveHalvingPruner(min_resource=WORKLOAD_SIZED[0], reduction_factor=10),
        sampler=RandomSampler(),
    )
    study.optimize(objective, timeout = STUDY_TIMEOUT)

    pruned_trials = [t for t in study.trials if t.state == optuna.structs.TrialState.PRUNED]
    timeout_trials = [t for t in study.trials if t.state == optuna.structs.TrialState.FAIL]
    complete_trials = [t for t in study.trials if t.state == optuna.structs.TrialState.COMPLETE]
    print('Study statistics: ')
    print('  Number of finished trials: ', len(study.trials))
    print('  Number of pruned trials: ', len(pruned_trials))
    print('  Number of timeout trials: ', len(timeout_trials))
    print('  Number of complete trials: ', len(complete_trials))

    print('Best trial:')
    trial = study.best_trial
    print('  Value: {}'.format(trial.value))
    print('  Params: ')
    for key, value in trial.params.items():
        print('    {}: {}'.format(key, value))