[Draft] 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と呼ばれています。

// TODO(tsuzuku) 図を足す。

Write-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を用いることで実現されています。

Read-write conflictの回避

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

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

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

Sampling algorithmの整理

Trialを[serializable](https://en.wikipedia.org/wiki/Isolation_(database_systems)にすることは困難なので、 Snapshot Isolationだけ要求することにします。 すなわち、”各trialはtrial開始時のhistoryの状態にのみアクセスできる”という条件を足します。 これはasynchronous algorithmでは珍しくない設定な気がします(解析とかしやすい)。

内部実装の話

Storage classでsnapshotをとっています。 DBにtimestampがあるべきでは?というのが最もな話なんですが、サポートしているDBが広くて難しかったです。

余談

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