Make requeue_zombie_trials() work with Optuna 2 and 3

mala-project · May 30, 2024 · e15c8e9 · e15c8e9
1 parent f7abbfc
commit e15c8e9
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/mala/network/hyper_opt_optuna.py b/mala/network/hyper_opt_optuna.py
@@ -176,9 +176,18 @@ def requeue_zombie_trials(study_name, rdb_storage):
         cleaned_trials = []
         for trial in trials:
             if trial.state == optuna.trial.TrialState.RUNNING:
-                study_to_clean._storage.set_trial_state(
-                    trial._trial_id, optuna.trial.TrialState.WAITING
+                kwds = dict(
+                    trial_id=trial._trial_id,
+                    state=optuna.trial.TrialState.WAITING,
                 )
+                if hasattr(study_to_clean._storage, "set_trial_state"):
+                    # Optuna 2.x
+                    study_to_clean._storage.set_trial_state(**kwds)
+                else:
+                    # Optuna 3.x
+                    study_to_clean._storage.set_trial_state_values(
+                        values=None, **kwds
+                    )
                 cleaned_trials.append(trial.number)
         printout("Cleaned trials: ", cleaned_trials, min_verbosity=0)