はじめに

みなさん、こんにちは。Optunaコミッタの今村です。OptunaではRDBサーバを利用することで複数台のマシンを使った分散最適化をサポートしています。一方で、数千を超えるプロセスからRDBサーバにアクセスするとサーバに対して大きな負荷がかかってしまい、速度低下を招く可能性があります。最新のOptuna v4.2.0リリースでは、このような大規模分散最適化によってRDBサーバが高負荷になるケースに対応するために、gRPCストレージプロキシと呼ばれる機能を導入しました。

本稿では、Optunaを利用した分散最適化の方法を説明し、大規模な分散最適化を実現するためのgRPCストレージプロキシをご紹介します。

TL;DR

• OptunaではRDBサーバを用いることで分散最適化が簡単に実現できます。
• gRPCストレージプロキシを利用することでRDBサーバの負荷を軽減し大規模な分散最適化を実現することができるので、ぜひ利用してみてください。

RDBサーバを利用した分散最適化

以下のような簡単なOptunaのコードを考え、分散最適化を実行してみましょう。このコードは単一のプログラムとして実行することもできますが、コードを全く書き換えることなく分散最適化に移行することができます。

 
import optuna

def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return x ** 2

if __name__ == "__main__":
    study = optuna.create_study(
        study_name="distributed-example", 
        storage="mysql://root@localhost/example",
        load_if_exists=True,
    )
    study.optimize(objective, n_trials=100)

Optunaを用いた分散最適化を実行するためには、最適化履歴を複数のプロセスから参照できるようにRDBサーバを利用する必要があります。上記のコードではMySQLサーバの接続URLがoptuna.create_study 関数のstorage引数に渡されているので、準備として以下のようにローカルのPCにMySQLを利用したRDBサーバを起動しておきましょう。インストールガイドを参考に、必要に応じてMySQLをインストールしてください。

mysql -u root -e "CREATE DATABASE IF NOT EXISTS example"

分散最適化の実行方法はとても単純です。PC上で複数のターミナルを開いてください。そして、各ターミナルで上記のコードを全く同じコマンドで実行するだけです。例えばコードが foo.pyに書かれてあるなら、各ターミナルでpython foo.pyを実行すれば良いです。RDBサーバに保存されたデータを各プロセスが参照することで、単一のプログラムと全く同じコードで分散最適化が実現できます。以下は実際に二つのターミナルで実行したコマンドの出力結果です。

• ターミナル1

$ python foo.py
[I xxx] Trial 0 finished with value: 45.35553104173011 and parameters: {'x': 8.73465151598285}. Best is trial 0 with value: 45.35553104173011.
[I xxx] Trial 2 finished with value: 4.6002397305938905 and parameters: {'x': 4.144816945707463}. Best is trial 1 with value: 0.028194513284051464.
...

• ターミナル2

$ python foo.py
[I xxx] Trial 1 finished with value: 0.028194513284051464 and parameters: {'x': 1.8320877810162361}. Best is trial 1 with value: 0.028194513284051464.
[I xxx] Trial 3 finished with value: 24.45966755098074 and parameters: {'x': 6.945671597566982}. Best is trial 1 with value: 0.028194513284051464.
...

gRPCストレージプロキシとは

RDBサーバさえ用意すれば簡単にOptunaでは分散最適化を利用することができます。しかし、Optunaの適用範囲が機械学習のハイパーパラメータ最適化から材料探索等に広がっていくにつれて、大規模かつサーバ処理負荷の大きな最適化が行われるようになりました。例えば、Preferred Networks社で行っている材料探索シミュレーションへの応用では、1 Trialに数分かかる設定で数万から数十万Trialを実行するStudyの最適化を数十並列で数百Studyも回しています。結果として並列数は数千ワーカーにもなり、RDBサーバに大きな負荷がかかり、クエリ応答を待つ時間がボトルネックとなっていました。

gRPCストレージプロキシは、Optuna v4.2.0で導入された大規模な分散最適化向けの機能です。具体的には、最適化実行ワーカーとRDBサーバの間に入り、Storage APIの呼び出しを中継します。数百から数千のワーカーが動作する大規模な分散最適化の設定では、数十台ごとにgRPCストレージプロキシを設置することで単一障害点となるRDBサーバの負荷を大幅に軽減できます。またOptunaのStudyやTrialといった情報をワーカーから共通に利用できるキャッシュとして管理することでも負荷の軽減が期待できます。

図：gRPCストレージプロキシのイメージ図

利用方法とベンチマーク

gRPCストレージプロキシを利用して分散最適化を行うには、上記のRDBサーバに加えて、最適化実行ワーカーとの間に入るプロキシサーバを立てておく必要があります。プロキシサーバを立てる機能はoptuna.storages.run_grpc_proxy_serverとして提供されています。以下は上記のMySQL RDBサーバに対して、localhostのポート番号13000番にプロキシサーバを立てるコードです。

import optuna


if __name__ == "__main__":
    optuna.storages.run_grpc_proxy_server(
        optuna.storages.get_storage("mysql://root@localhost/example"),
        host="localhost",
        port=13000,
    )

上記のプロキシサーバを経由してRDBサーバを利用するためには、最適化を実行するワーカーのコードを修正する必要があります。具体的にはoptuna.create_study関数のstorage引数にoptuna.storages.GrpcStorageProxy を渡す必要があります。下記のコードでは、プロキシサーバを指定する際に用いたホスト名localhostとポート番号13000番を用いています。

import optuna


def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return x ** 2


if __name__ == "__main__":
    study = optuna.create_study(
        study_name="distributed-example",
        storage=optuna.storages.GrpcStorageProxy(host="localhost", port=13000),
        load_if_exists=True,
    )
    study.optimize(objective, n_trials=100)

ポート番号を変えることで異なるプロキシサーバを経由するようにできますし、また異なるホスト名を指定することでネットワークを跨いでプロキシサーバを利用することもできます。例えば、Kubernetes等によって管理された複数ノードの環境では、各ノード上のPodに一つのプロキシサーバを立て、Pod内の複数のContainerからPodのプロキシサーバを経由する構成にすることで、別のPodに立てたRDBサーバへの負荷を軽減しながら分散最適化を実行することもできます。

gRPCストレージプロキシの概念実証のため、以下のような設定で性能比較を行いました。300台のワーカーに対して10台ごとにプロキシサーバを設置しRDBサーバへの通信を中継することで、10000 Trialsの分散最適化をRandomSamplerを用いて実行しました。以下はワーカーごとにかかった処理時間の平均値です。この実験シナリオでは処理時間を6割削減できました。

表：大規模分散最適化時の速度改善

なお、今回はMySQLサーバに大きく負荷がかかる設定となっており、MySQLサーバが十分に処理できるワーカー数であればgRPCのオーバーヘッドの分だけgRPC Storage Proxyの方が遅くなる可能性があります。実際、単一ワーカーでは1.2~1.3倍程度の速度低下を確認しています。あくまでgRPCストレージプロキシは大規模な分散最適化向けの機能であることにご注意ください。

おわりに

Optunaを利用すれば分散最適化を簡単に実現することができます。また大規模な分散最適化を行う場合はgRPCストレージプロキシを利用することで、RDBサーバの負荷を軽減することができます。我々はこれからも大規模な分散最適化に向けた機能を拡張していきます。ぜひOptunaの分散最適化とgRPCストレージプロキシを利用してみてください。

それではよいブラックボックス最適化ライフを！