The Internal of Optuna’s Jupyter Lab and VS Code Extensions
Introduction
As we announced in the Optuna 3.4 release blog, we have introduced Optuna Jupyter Lab extension and the VS Code extension. These extensions enable easy access to the Optuna Dashboard within Jupyter Lab and VS Code, making it convenient for you to view your optimization history, hyperparameter importances, etc. in graphs and tables.
This article aims to elaborate on the process of implementing these extensions and their functioning mechanism. Our hope is that this explanation will provide beneficial information for anyone interested in developing extensions for Jupyter Lab or VS Code.
Optuna Dashboard consists of a server program written in Bottle web framework and a single page application written in React, as shown in the following figure.
To integrate the Optuna Dashboard into Jupyter Lab, several significant modifications were necessary:
- Adapting the Optuna Dashboard’s server, originally written in the Bottle framework, to run within Jupyter Lab.
- Substituting the JavaScript HTTP client library with the @jupyterlab/services/ServerConnection.
- Developing a Jupyter Lab-specific user interface (including input screens for Optuna’s storage URL, etc.).
- Fixing the UI which is broken in Jupyter Lab (e.g. correcting display cut off)
Though these tasks were time-consuming and demanding, this article will primarily focus on how to operate a Bottle application within Jupyter Lab. Given that Jupyter Lab employs the Tornado web framework, its extension must be implemented as a Tornado handler like this official example. Consequently, running the Optuna Dashboard server within Jupyter Lab is not a straightforward endeavor.
Reimplementing all of the Optuna Dashboard server’s procedures from Bottle to Tornado would not have been a feasible solution, considering the implementation effort and subsequent maintenance costs. Therefore, we decided to invoke the processing of the Bottle application on the Tornado handler by converting the request object coming to the Tornado handler into a form comprehensible to Bottle. When we call procedures written in Bottle from the Tornado handler, there are two major options to consider.
- Calling WSGI entrypoint from the Tornado handler
- Calling Bottle’s view functions directly from the Tornado handler
In conclusion, we implemented this with the former approach. We avoided the latter approach because the request object in the Bottle framework is defined as a variable which is thread-local and module global, and they can be imported and utilized by view functions. Invoking Bottle’s view functions from within a Tornado handler requires the initial setup of request object and response object. Compared to the former approach, the latter approach had more points to consider and was thought to be more complex.
The former approach is not overly complex given you have an understanding of the WSGI specifications (refer to PEP 333 and PEP 3333 for more information about WSGI). After converting the request object passed to the Tornado handler into a WSGI environment dictionary object, you can invoke the WSGI entrypoint. This is achieved by mapping the status code and HTTP headers set in the start_response() function and the response body obtained from the function return value to the Tornado response object. Originally, the plan was to implement these processes from scratch, but fortunately, suitable processes were already established in the tornado.wsgi module.
A potential concern is the fact that Tornado is an asynchronous web framework, while the WSGI application is just a function (callable object). This discrepancy may block the Tornado server’s eventloop for prolonged periods while waiting for a response from the WSGI application. If such practical issues arise, one possible workaround would be to use multi-thread to prevent the eventloop from blocking. While I have not yet conducted a detailed investigation on this, if you encounter any issues such as Jupyter Lab becoming slow while using the Optuna Dashboard extension, I would appreciate your feedback.
VS Code Extension
The process of Implementing the VS Code extension is more complex than that of the Jupyter Lab extension, primarily because it fundamentally relies on Node.js. Consequently, we have decided to depart from the full functionality support of the Optuna Dashboard, as executed in the Jupyter Lab extension. Our plan is to create an application that leverages WebAssembly, removing the need for a Python web server. This application will then be encapsulated using WebView and offered as a VS Code extension.
The independence from Python, coupled with the ability to run on a web browser, has made it possible to host the application on GitHub Pages. Basic functionalities such as visualizing optimization history and the importance of hyperparameters are already supported. Thus it’s ready-to-use when you need a swift review of the optimization history.
The most crucial responsibility of the Python web server is loading Optuna studies and trials from the storage. It poses a significant challenge to support databases like MySQL and PostgreSQL — which transmit data via a TCP connection — to run on a web browser. Consequently, the strategy for the VS Code extension was to restrict support to file-based storages such as SQLite3 and JournalFileStorage.
We utilize @sqlite.org/sqlite-wasm for data loading from SQLite3. Optuna’s database manages the schema with the DB migration tool Alembic, and changes are made as needed through migration. Given that the table definition may vary based on the Optuna version in use, we issue suitable SQL statements depending on the schema version. This complexity in schema handling presents difficulties for utilizing solutions such as ORM. Consequently, in the VS Code extension, we assemble and issue all raw SQL.
Currently, it only supports SQLite3, but it would be possible to support JournalFileStorage by implementing a process to restore the state from operational logs saved in JSON Lines format. We might be able to prioritize this if we receive feedback through GitHub issues or Visual Studio Marketplace reviews. We look forward to your feedback.
Furthermore, the VS Code extension computes hyperparameter importances directly on the web browser leveraging WebAssembly. The task of implementing fANOVA (an algorithm used to calculate hyperparameter importance) from scratch is challenging due to its requirement for a random forest. Fortunately, Takeru Ohta (@sile), an Optuna committer, had already implemented fANOVA and random forest in Rust. I prepared a JavaScript interface using wasm-bindgen and utilized it in the extension.
The current release of the VS Code extension does have some areas of implementation that could be improved upon. In future updates, we plan to focus on the following points. If you have any requests or suggestions, we welcome them via GitHub issues or pull requests.
- Support for JournalFileStorage
- Compatibility with code-server
- Reflect changes in storage onto the Dashboard screen using the Custom Editor API
- Perform data loading from storage on the Web Worker.
Conclusion
In this article, we explained how the Optuna Dashboard is run within Jupyter Lab and VS Code, and the mechanism behind it. If this article has piqued your interest in the development of Optuna Dashboard, we would be delighted if you would join the development, referring to this document.
In conclusion, I need to underscore that these achievements are the result of collective efforts. The development of the Jupyter Lab extension was aided by my collaboration with Shinichi Hemmi (@Alnusjaponica), an Optuna committer. Moreover, Yuiga Wada (@YuigaWada) contributed to supporting Optuna’s older database schema, and @ciffelia contributed to enhancing error handling in the Rust program. I would like to express my gratitude for everyone’s invaluable contributions to the Optuna Dashboard.
