Optuna not parallelising across workers using joblib spark backend on Databricks

[alexisdrakopoulos]

Code:

import optuna
import joblib
from joblibspark import register_spark
import mlflow

def objective(trial):
    from time import sleep
    sleep(60)
    return 1

mlflow.autolog(disable=True)

study = optuna.create_study(direction="maximize")

register_spark()
 # register Spark backend for Joblib
with joblib.parallel_backend("spark", n_jobs=-1):
    study.optimize(objective, n_trials=100)

Using LTS DBR ML 14.3. Only extra library installed is !pip install Optuna.

Our machine has 8 active workers and I am running this on Databricks on the driver.

It clearly runs 1 trial at a time and does not parallelize them. I tried following all the online docs but can't see anything. No warnings or anything appear:
[I 2025-03-18 15:12:04,802] A new study created in memory with name: no-name-ca4eddd4-26b3-4e78-b611-849d60291c94
[I 2025-03-18 15:13:04,879] Trial 0 finished with value: 1.0 and parameters: {}. Best is trial 0 with value: 1.0.
[I 2025-03-18 15:14:04,933] Trial 1 finished with value: 1.0 and parameters: {}. Best is trial 0 with value: 1.0.

Any advice is appreciated as it's critical to us.

I also tried LTS 12.2 and 15.3 but no changes.