Prototyping a Faster Optuna Implementation in Rust
Introduction
Optuna is implemented entirely in Python without relying on C extensions. It enjoys smooth installation across various execution environments — a critical attribute given its widespread usage. At the moment, we have no plans to depend on non-Python languages.
On the other hand, there have been several use cases where execution speed matters, or where we would like to integrate existing systems coded by another language. With the recent notable speed enhancements in Python toolchains, especially Ruff which uses Rust, we have started closed-source prototyping of Optuna implementation in Rust. While we are not currently planning for a public release, this blog outlines the motivation behind this development and the potential benefits that a Rust-based Optuna implementation could provide.
Motivation
The Optuna implementation in Rust is primarily motivated by two factors.
Making Optuna Faster
Optuna primarily targets hyperparameter optimization in machine learning. In such scenarios, model training and evaluation are typically time-consuming processes, so Optuna’s execution time seldom becomes the bottleneck.
However, there are many possible applications of black-box optimization beyond the optimization of machine learning hyperparameters. For example, in the case presented below where Optuna is used in the field of material science to explore unknown crystal structures, the execution speed of Optuna becomes critical because many trials are required. The more iterations conducted, the higher the probability of finding unknown crystal structures.
Efficient Crystal Structure Prediction using Universal Neural Network Potential and Genetic Algorithm
Broadening Language Support
The Rust-based implementation can be used not only from Rust and Python, but also provides JavaScript bindings and a C-API. In particular, we aim to utilize the JavaScript bindings in Optuna Dashboard.
Optuna Dashboard is a web dashboard used to analyze Optuna trial results. We also recently started providing extensions for Jupyter Lab and VS Code.
While implementing these, we have re-implemented various logics in TypeScript, including the visualization module (optuna.visualization) and the JournalFileStorage loader implemented in Optuna. Some of these have speed issues, and it is expected that the Rust-based implementation can improve the user experience by providing fast JavaScript bindings. There are also development benefits, which we will detail later.
Accelerating Optuna Using Rust-based Implementation
In our prototype of the Rust-based Optuna implementation, we have re-implemented the core functionality of Optuna in Rust while also facilitating Python bindings using PyO3. This section elaborates on the implementation policy of the Rust-based Optuna implementation and its execution speed as of now.
Compatibility and Speed Comparison with Optuna
As Ruff provides high compatibility with existing toolchains, the Rust-based Optuna implementation also emphasizes compatibility with the Optuna API. Users don’t need to learn a new API, they can straightforwardly transition from existing projects and continue to reap the benefits of the ecosystem that Optuna has built. Since the Rust-based Optuna implementation provides essentially the same API as Optuna, the following code works as is, by just changing the import statement.
import <OPTUNA_RUST> as optuna
def objective(trial:optuna.Trial) -> float:
x = trial.suggest_float('x', -10, 10)
y = trial.suggest_float('y', -10, 10)
value = (x - 2) ** 2 + (y + 5) ** 2
return value
study = optuna.create_study()
study.optimize(objective, n_trials=1000)
print(study.best_trial)The current performance comparison for the above code, when varying the number of trials and parameters, is as follows.
Optuna uses TPESampler by default. As the comparison table in the official documentation suggests a recommended budget of 100–1000 for the TPESampler, the TPE is not necessarily the suitable method when many trials can be evaluated from the perspectives of computational complexity and optimization performance. If we have an abundant computational budget, other methods are often recommended. For example, the comparison of execution times for the RandomSampler is as follows.
One key observation from the above results is that when we increase the number of trials tenfold, the execution time of the Rust-based implementation only takes about ten times longer. In contrast, the execution time of Optuna increases by 50 to 100 times. This can be attributed to some of the convenient features provided by Optuna. At Optuna, we carefully balance convenience and speed reduction in our development, but it is also true that as a result of pursuing convenience for many users, we have caused such slow-down problems for some users.
The Rust-based implementation puts a high priority on speed. If there isn’t a significant speed improvement, the Rust-based implementation is not necessary. If we heavily focus on compatibility with the Optuna design and its functionalities, Rust-based implementation may not be able to achieve significant performance improvements. Hence, we are carefully considering the implementation of certain features, such as enqueue_trial(). This is one of the reasons why I mentioned that the Rust-based implementation provides “essentially” the same API.
Integration of Optuna (Python) and Its Rust-Based Implementation
The Rust-based implementation is not aimed to serve as a complete replacement for Optuna. Instead of re-engineering all of Optuna’s features in Rust, we aim to create a cooperative library where the strengths of each can be utilized and complement each other.
For instance, RDBStorage is an important feature in Optuna. It enables the persistence of trial results and distributed optimization. However, re-implementing this feature in Rust presents only a minimal advantage. The main bottleneck lies in I/O operations, which don’t typically benefit from re-implementation in terms of speed. Consequently, the Rust implementation of Optuna offers conversion utilities to avoid redeveloping such features.
import <OPTUNA_RUST> as optuna
from <OPTUNA_RUST> import FromOptunaStorage
from optuna.storages import RDBStorage
def objective(trial: optuna.Trial) -> float:
...
storage = FromOptunaStorage(RDBStorage("sqlite://db.sqlite3"))
study = optuna.create_study(storage=storage, study_name="example")
study.optimize(objective, n_trials=10)In the example above, we are partly leveraging Optuna’s capabilities from the Rust-based implementation. The opposite is also possible. The following code illustrates partially utilizing the Rust-based implementation within Optuna.
import optuna
# from optuna.importance import FanovaImportanceEvaluator
from <OPTUNA_RUST>.optuna.importance import FanovaImportanceEvaluator
def objective(trial: optuna.Trial) -> float:
...
study = optuna.create_study()
study.optimize(objective)
importance = optuna.importance.get_param_importances(
study, evaluator=FanovaImportanceEvaluator()
)Though this approach has some limitations in terms of speed improvement compared to the previous one, it allows you to try the Rust-based implementation more easily without significantly rewriting existing code. If any problem arises or the performance advantages aren’t realized, it’s much easier to switch back.
Handling Technical Challenges in Optuna Dashboard
Lastly, I will introduce the vision for utilization in the Optuna Dashboard. The Optuna Dashboard is a web application for viewing Optuna trial results. The development of this web dashboard involves two major technical challenges: improving execution speed and streamlining the process of unit testing.
Making Optuna Dashboard Faster
The Optuna Dashboard consists of a Python API server and a Single Page Application (SPA). The main role of the API server is to just deliver Studies and Trials to the web frontend via a JSON API so that most of the important application logics are implemented in the single page application using TypeScript. This greatly contributes to the simplicity of the design of Optuna Dashboard and the improvement of user experience, and is also very useful for providing the VS Code extension.
However, unfortunately, the current Optuna Dashboard is quite slow and hard to use when dealing with tens of thousands of trials. There are several fundamental factors for this, but if the Rust-based implementation provides fast JavaScript bindings, Optuna Dashboard will become faster while maintaining the current simple design and improving the user experience.
Writing Maintainable Unit Tests
Most of the frontend logics of the Optuna Dashboard takes a list of Optuna Trials as arguments. Trials contain a lot of necessary information, such as params and distributions, which can make the setup code wordy and complicated when generating multiple trials. Introducing utilities for generating fixtures can be considered, but there are concerns that this itself can become bloated and create a debt of not knowing what is being tested.
By providing JavaScript bindings for the Rust-based implementation, we can easily generate Study or Trial information on the Node.js side for testing. There is no need to hardcode distributions or internal_params etc. to prepare Trials, making it strong against changes in internal structure and improving the maintainability of unit tests.
import * as optuna from '<OPTUNA_RUST>.js';
const study = optuna.create_study();
const objective = (trial) => {
const x = trial.suggest_float("x", -10.0, 10.0);
const y = trial.suggest_int("y", -10, 10);
const z = trial.suggest_categorical("z", ["foo", "bar"]);
return value = (x - 5) ** 2 + (y + 2) ** 2
}
study.optimize(objective, 10)
console.log(study.best_trial())Until now, development in Optuna Dashboard has heavily depended on E2E testing, executing Optuna, saving trial results into SQLite3, launching the actual Python API server, and validating the display with pytest and Playwright. While E2E testing allows us to test Optuna Dashboard easily, this approach is not suitable to increase the test coverage. Providing the JavaScript bindings allows us to write maintainable unit tests, leading to improved test coverage.
Conclusion
In this article, I introduced the motivation and vision for the utilization of the newly initiated Rust-based Optuna implementation. As mentioned at the beginning, the Rust-based Optuna implementation is still in the prototype stage, and there are no specific plans for public release at this time. We plan to continue developing and verifying its usefulness within the committers.
