Code for reproduction for Optuna PR #1498.

pr1498.py

from multiprocessing import Pool
import os
import sys

import optuna


def f(x, y):
    return (x - 3) ** 2 + y


# Objective function
def optuna_objective(trial):
    x = trial.suggest_uniform("x", -10, 10)
    y = trial.suggest_uniform("y", -10, 10)
    z = f(x, y)
    return z


# Entry point
def main(args):
    (rank, study_name) = args
    # Create a study
    study = optuna.create_study(
        study_name=study_name,
        storage=os.environ["OPTUNA_SQL"],  # Set a remote PostgreSQL server
        load_if_exists=True,
        direction="minimize",
    )
    # Run optimization
    study.optimize(optuna_objective, n_trials=20)


if __name__ == "__main__":
    study_name = sys.argv[1]
    n_workers = int(sys.argv[2])
    filename = sys.argv[3]
    with Pool(n_workers) as pool:
        pool.map(main, [(i, study_name) for i in range(n_workers)])

    study = optuna.create_study(
        study_name=study_name,
        storage=os.environ["OPTUNA_SQL"],  # Set a remote PostgreSQL server
        load_if_exists=True,
        direction="minimize",
    )
    study.trials_dataframe().to_csv(filename)