Optuna parameter optimisation with MPI

Question

I have some machine learning code which uses SVM (from scikit-learn) with a pre-computed kernel that I want to optimise using optuna, so the code simplistically looks a bit like this

def objective(trial: Trial, fast_check=True, target_meter=0, return_info=False):
     #Set parameters
     C = trial.suggest_float("C",0.01,5)
     tol = trial.suggest_loguniform("tol",1e-4,1e-1)
     kernel params = ...

     #Build train kernel
     kernel_train = ...

     #Build test kernel
     kernel_test = ...

     #Train svc
     svc = SVC(kernel="precomputed", C=C, tol=tol)
     svc.fit(kernel_train, train_labels)
     test_predict = svc.predict(kernel_test)
     test_auc = roc_auc_score(test_labels,test_predict)

     return test_auc

study = optuna.create_study(direction="maximize",study_name='study_1')
study.optimize(objective, n_trials=40)

However due to the complicated nature of the kernel I am computing, I am using mpi4py to parallelise the calculations but I am having some issues using both optuna and MPI together.

Obviously I want the kernel code over the multiple processors but when I create the study and optimise it, I don't want to then create multiple different studies across the processors I just want one study on root that is optimised (I assume?). I have tried the below and it works however it doesn't optimise as well as when I am not using MPI and I think this is creating multiple studies and optimising them all which doesn't seem efficient. It seems to struggle more to converge on the optimal parameters.

from mpi4py import MPI

mpi_comm = MPI.COMM_WORLD
rank = mpi_comm.Get_rank()
n_procs = mpi_comm.Get_size()
root = 0

def objective(trial: Trial, fast_check=True, target_meter=0, return_info=False):
     #Set parameters
     C = trial.suggest_float("C",0.01,5)
     tol = trial.suggest_loguniform("tol",1e-4,1e-1)
     kernel params = ...

     #Build train kernel using MPI
     kernel_train = ...

     #Build test kernel using MPI
     kernel_test = ...

     #Train svc
     if rank == root:
           svc = SVC(kernel="precomputed", C=C, tol=tol)
           svc.fit(kernel_train, train_labels)
           test_predict = svc.predict(kernel_test)
           test_auc = roc_auc_score(test_labels,test_predict)
     else:
           test_auc = 0
     test_auc = mpi_comm.bcast(test_auc, root=0)

if rank == root: 
     study = optuna.create_study(direction="maximize",study_name='study_1')
else:
     study = 0
 study = mpi_comm.bcast(study, root=0)

study.optimize(objective, n_trials=40)

It's a very niche question but just wondering if anyone has any experience with either of these packages and could help suggest how to run a multi-processing code while only optimising the parameters on one processor. Apologies if any of the terminology isn't correct, I am new to using both of these packages so some patience would be appreciated. :)