Performance Improvement

This change improves performance by reducing the overhead of fetching all trials. Below are benchmark results comparing the current master branch.

Before (master)

After (PR)

Here is the code used to generate the results above

import optuna
import cProfile

study = optuna.create_study()

def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    y = trial.suggest_uniform("y", -10, 10)
    return x**2 + y**2

study.optimize(objective, n_trials=100)

for i in range(-5, 5):
    study.enqueue_trial({"x": i, "y": i})

def profile_objective():
    study.optimize(objective, n_trials=100)

cProfile.run('profile_objective()', 'journal_storage_master_100.prof')
# cProfile.run('profile_objective()', 'journal_storage_fix_100.prof')

@not522 Could you review this PR?

Let me show you the performance improvement with various examples.

Performance Improvement

Below are benchmark results comparing the performance before and after this PR using different scenarios.

1. study.optimize() for n_trials = 100, 1000

This PR significantly improves performance because _get_all_trials(state=TrialState.WAITING) now only checks the last trial's state instead of iterating over all trials.

Here is the code used to generate this data.

import optuna
import cProfile

study = optuna.create_study()

def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    return x**2

def profile_objective():
    study.optimize(objective, n_trials=num_of_trials)

profiller = cProfile.Profile()
profiller.run('profile_objective()')

2. study.enqueue_trial() followed by study.optimize()

I was not really sure why this test generates some difference.

Here is the code used to generate this data.

import optuna
import cProfile

study = optuna.create_study()

def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    return x**2

for i in range(-num_of_waiting_trials // 2, num_of_waiting_trials // 2):
    study.enqueue_trial({"x": i})

def profile_objective():
    study.optimize(objective, n_trials=1)

profiller = cProfile.Profile()
profiller.run('profile_objective()')

3. study.optimize() for n_trials = 500, followed by study.enqueue_trial() for 100 trials, then study.optimize(n_trials=1), profiled with cProfile

The performance improvement comes from optimizing the logic: previously, the function checked all 500 + 100 trials, whereas the new implementation only checks 100 trials.

Here is the code used to generate this data.

import optuna
import cProfile

study = optuna.create_study()

def objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    return x**2

study.optimize(objective, n_trials=500)

for i in range(-50, 50):
    study.enqueue_trial({"x": i})

def profile_objective():
    study.optimize(objective, n_trials=1)

profiller = cProfile.Profile()
profiller.run('profile_objective()')

Benchmarking Details

• Each test, except for the one with n_trials = 10,000, was conducted 10 times, alternating between the master branch and the PR branch in every iteration (i.e., master → PR → master → PR → ...).
• The average and standard deviation were computed.
• Raw data is available here.

Could you share benchmarks with larger n_trials (e.g., 10,000 or 50,000)? When it is small, _pop_waiting_trial_id is fast enough to be ignored.

Sure. But let me do it with n_trials = 10,000 and for a single iteration since my PC isn't powerful enough and doing optimization with larger number of n_trials needs significant amount of time

@not522 I added the data to the comment above. The improvement is significant—total execution time for pop_waiting_trials (10,000 runs) dropped from 24.929083 seconds to just 0.391722 seconds.