Optunaの分散実行方式

Optunaの分散実行について書こうと思います。 WriterはOptuna committer兼Optunaの開発を行なっているPreferred Networks, Inc.の社員です(コンプラ重視)。 尚、本記事は非公式documentで、optuna teamやPreferred Networks社を代表するものではありません。 想定読者を限定しませんが、committer以外には非自明な仮定を用いているかもしれません。

構成

Optunaの分散最適化は、1つのdatabaseと複数のworkerで構成されます。

実行する計算

過去の最適化の履歴全体を受け取って、次に試すべきパラメータをsuggest/evaluateし、それを(databaseに)feedbackする、という操作を繰り返します。 理想的には上の一つの操作を一つのtransactionとして実行できるといいのですが、それだとserial実行しかできないので最終的には緩和した問題を考えます。 これについては後述します。 Optunaでは上述のsuggestからfeedbackまでの操作一つがTrialと呼ばれ、管理されています。 まあた、関連するtrialの集合がstudyと呼ばれています。

Write conflictの回避

transactionについて考える前にどうやってwriteのconflictを回避しているかについて書きます。

Study

Study(のmetadata)へは書き込み自体が稀で、起きた場合の挙動はundefinedです。

Trials

Trialへの書き込み競合の回避はtrialをRUNNINGにしたworkerしか書き込めない、という制約とそのようなworkerは一つしか存在できない、という制約によって実現されています。 このRUNNING stateはlockと見なすことができます。 このlockを獲得する方法は現在二種類あります。

  • RUNNING stateのtrialを無から作る
  • WAITING stateのtrialをRUNNINGにする 後者を行う場合にはtest-and-setと同じ仕組みlockを取得します。 test-and-setの機構自体はDBのlockを用いることで実現されています。

Repeatable-read

実は僕が入社した時点では書き込みの競合までしかcareされていなかったのですが、最近のoptunaではそれだと問題になるケースがあるのでもう少し話を進めます。

どういう問題が残っているか

上述の書き込みは実はexclusive-lockにはなっていません。 というのも、他のtrialが書き込みをしたりstate変更をしている最中でもそのtrialを他から読めるからです(dirty read)。 これによって引き起こされる大きな問題は、trialの実行中にtrialに渡される過去のhistoryが刻一刻と変わってしまうことにあります。 これは基本的にsampling algorithmの実装では考慮されていないので、良くてperformance低下、悪いと任意のerrorで異常終了します。 例えば、sampler classが一つのtrialの中使える計算結果をcacheしようとした時、この辺りが議論されていないと壊れます。 ここからはどうこの問題を解決しているかについて話していきます。 この問題発見と続くfixは僕のonboarding taskの一環でした。

Sampling algorithmの変更

見通しが良いので、Trialをserializableにしていく、というのを目標に話を進めます。 これが満たされてれば上述の問題は消えます(repeatable-read)。 尚、Commitをする、という操作はtrialを終了statusにする、とう操作と対応しています。 まず、上述のすべてのtransactionをserializableにしようとするとserialな実行しかできないので、問題を緩和する必要があります。 具体的には、”各trialはtrial開始時のhistoryの状態にのみアクセスできる”という条件を足し、そしてtrialがserializableであるという条件を外します。 これはasynchronous algorithmでは珍しくない緩和な気がしています(解析とかしやすい)。

内部実装の話

実装としてはSnapshot Isolationでrepeatable-readにしています。 といっても、database server側でこれはできていなくて、Storage classが代わりにsnapshotをとっています。 尚、前述のlockのおかげでwrite-write conflictは起こりません。

余談

  • optuna開発から離れて久しいので微妙に仕様が変わっていたりするかもしれません。
  • 最近複数workerがstateをRUNNIINGにできるようにするという話があるらしいので気になっています。今回これを書こうと思ったきっかけもそれです。
  • 本当はsnapshotはstorage classではなくDB + sampler/prunerでsupportするべきで、そうしていないせいで実はまだ壊れるケースとかが放置されています(がsamplerなどは互換性の問題でなかなか手を入れられない)。
  • すべてのstorageにsnapshotを入れる話まではあったと思うのですが、一方で今はrdb backend以外の対応は済んでいなかったと思うので、分散最適化中implicitに壊れうるなという気がしています。