Blog
はじめに
Optunaの新しいバージョン3.4では、新たにJupyter Lab拡張とVS Code拡張が公開・アナウンスされました。これらの拡張機能を利用することで、Optuna DashboardをJupyter LabやVS Code内で起動し、最適化履歴をより手軽に確認できます。
本記事ではこれらの拡張機能がどのように実装されているのか、その仕組みを解説します。Optuna Dashboardの開発に興味を持ってくださっている方に限らず、Jupyter Lab拡張やVS Code拡張を開発しようとしている方のお役に立てば幸いです。
Jupyter Lab拡張の仕組み
Optuna Dashboardは次の図に示すように、PythonのBottleフレームワークで書かれたサーバープログラムと、Reactで書かれたシングルページアプリケーションからなります。
このプログラムをJupyter Lab拡張として動作するように書き換える際には、主に次のような作業が必要となりました。
- Bottleフレームワークで書かれたサーバープログラムを、Jupyter Lab内で動かす
- JavaScriptのHTTPクライアントをJupyter Labが提供する @jupyterlab/services/ServerConnection を使った処理に差し替える
- Jupyter Lab用のUI (OptunaのストレージURL入力画面等) を実装する
- 画面が見切れるなどの表示崩れを修正する
どれも少々手間のかかる作業ではありましたが、ここではその中でも技術的にやや特殊な対応が必要となった「Bottleで書かれたサーバープログラムをどのようにJupyter Labの中で動かすか」という点に絞って解説していきます。Jupyter LabはTornadoというWebフレームワークで実装されたソフトウェアであることから、拡張機能についても基本的にはこちらの公式ExampleのようにTornadoのHandlerとして実装する必要があります。
Optuna Dashboardサーバーの処理をすべてBottleからTornadoで再実装するのは、実装工数やその後のメンテナンスコストの観点からあまり現実的な選択肢ではありません。そこでTornado Handlerに渡ってくるリクエスト情報をBottleが理解できる形式に変換することで、Tornado Handler上でBottleアプリケーションの処理を呼び出すことにしました。Tornado HandlerからBottleで記述された処理を呼び出す際には、大きく次の2つの選択肢が考えられます。
- Tornado Handlerから、WSGIエントリーポイントを呼び出す
- Tornado Handlerから、BottleのView関数を直接呼び出す
結論からいうと今回は前者のアプローチで実装しました。後者のアプローチを避けた理由はBottleフレームワークではリクエストオブジェクトが スレッドローカルでモジュールグローバルな変数として定義されており、各ビュー関数がそれらをインポートして利用することがあるためです。Tornado Handler内からBottleのビュー関数を呼び出す際には、事前にリクエストオブジェクトやレスポンスオブジェクトを適切に初期化しておく必要がありますが、前者のアプローチに比べて気にするべき点が増え、複雑さが増すと考えました。
前者のアプローチは、WSGIの仕様を把握していればそれほど難しいものではありません（WSGIの詳細は PEP 333およびPEP 3333をご確認ください）。Tornado Handlerに渡ってくるリクエストオブジェクトをWSGI Environment辞書オブジェクトに変換してからWSGIエントリーポイントを呼び出し、start_response関数にセットされるステータスコードやヘッダー情報、関数の返り値から得られるレスポンスボディーをTornadoのレスポンスオブジェクトにマッピングすれば実装できます。当初はこれらの処理をフロムスクラッチで実装しようとしていましたが、幸いにも tornado.wsgiモジュールに同様の処理が実装されていることがわかり、こちらを利用して動かすことができました。
Tornado HandlerからBottleを呼び出すため、変換レイヤーを挟む
1つ懸念として考えられる点は、Tornadoが非同期のWebフレームワークであるのに対して、WSGIアプリケーションの呼び出しはただの関数呼び出しである点です。単純な実装ではWSGIアプリケーションがレスポンスを返すまでの間、Tornadoサーバーが管理しているイベントループを長時間占有してしまう可能性が考えられます。もし実用上の問題が見つかってくれば、View関数の呼び出しをマルチスレッドで行うなど、Tornadoサーバー側のイベントループを占有しない工夫が必要になるでしょう。
今回こちらに関する詳しい調査はまだしていませんが、Optuna Dashboard拡張を使用しているとJupyter Labの動作が遅くなるといった問題に遭遇した際は、ぜひ報告いただければ幸いです。
VS Code拡張の仕組み
VS Code拡張は原則としてNode.jsで記述する必要があることから、Jupyter Lab拡張ほど簡単ではありません。そこでJupyter Lab拡張のようにOptuna Dashboardの全機能をサポートすることは諦めることにしました。方針としてはWebブラウザ単体で動作するアプリケーションを作成し、それをWebViewでラップしてVS Code拡張として提供することにします。
Pythonの依存がなくWebブラウザ上で動作するため、GitHub Pagesでホスティングして提供することも可能になりました。最適化履歴やハイパーパラメーターの重要度の可視化など基本機能はすでにサポートされているため、履歴をサクッと見たいというときには手軽に利用できます。
GitHub PagesでホスティングされているWebブラウザ版Optuna Dashboard
https://optuna.github.io/optuna-dashboard/
Pythonのサーバープログラムが担っている最も重要な処理は、OptunaストレージからStudyやTrialを読み込むことです。Webブラウザ上で動作させるためMySQLやPostgreSQLのようにTCPのコネクション等を経由してデータを送受信するようなデータベースのサポートは困難です。VS Code拡張ではSQLite3やJournalFileStorageといったファイルベースのストレージに絞ってサポートすることにしました。
SQLite3からのデータ読み込みには@sqlite.org/sqlite-wasmを使用しています。OptunaのデータベースはDBマイグレーションツールAlembicを用いてスキーマを管理し、適宜マイグレーションにより変更が加えられています。Optunaのバージョンによりテーブル定義も異なることがあるため、スキーマのバージョンに応じて適切なSQLを発行しています。そのためORMのようなソリューションの採用は難しく、VS Code拡張ではすべて生のSQLを組み立て、発行しています。
現状ではSQLite3のみをサポートしていますが、JournalFileStorageもJSON Linesフォーマットで保存されている操作ログから状態を復元する処理を実装すれば対応が可能です。GitHub issuesやVisual Studio Marketplaceのレビュー等でフィードバックいただければ優先度をあげて対応できるかもしれません。みなさまからのフィードバックをお待ちしております。
その他にもVS Code拡張ではハイパーパラメーターの重要度をWebAssemblyを使ってブラウザ上で計算しています。ハイパーパラメーターの重要度を計算するfANOVAの実装にはランダムフォレストが必要になることから、フロムスクラッチでの実装は非常に大変です。幸いなことにOptunaのコミッターでもあるTakeru Ohtaさん(@sile)がRustでfANOVAやランダムフォレストを実装し、公開されていたので、wasm-bindgenを用いてJavaScript用のインターフェイスを用意し、利用させていただきました。
VS Code拡張の使い方
今回公開したVS Code拡張はまだまだ実装面で改善できる箇所が残っています。将来的には次のような点に対応していく予定です。他にもご要望やご提案をお持ちの方はGitHub Issueやプルリクエストをお待ちしております。
- JournalFileStorageをサポートする
- code-server拡張として利用可能にする
- Custom Editor APIを用いてストレージに変更があったときにDashboardの画面にも反映する
- ストレージからのデータ読込処理をWeb Worker側で行う
まとめ
本記事ではOptuna DashboardをどのようにJupyter LabやVS Codeの中で動かしているのか、その仕組みを解説しました。本記事を通してOptuna Dashboardの開発に興味を持ってくださった方がいれば、こちらのドキュメントを参考にしながら開発にご参加いただければ幸いです。
最後になりましたが、これらの成果は私一人によるものではなく、Jupyter Lab拡張はOptunaコミッターでもあるShinichi Hemmiさん (@Alnusjaponica) と協力して実装しました。またVS Code拡張の実装にあたって、Yuiga Wadaさん(@YuigaWada)はOptunaの古いデータベーススキーマのサポート、@ciffeliaさんはRustのコードのエラーハンドリングの改善を手伝っていただきました。みなさまのOptuna Dashboardへの貢献に感謝申し上げます。
