diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 812201ad..b91ce5f0 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -152,6 +152,7 @@ def create_app( storage=storage, study_name=dst_study_name, directions=src_study.directions ) dst_study.add_trials(src_study.get_trials(deepcopy=False)) + note.copy_notes(storage, src_study, dst_study) except DuplicatedStudyError: response.status = 400 # Bad request return {"reason": f"study_name={dst_study_name} is duplicaated"} diff --git a/optuna_dashboard/_note.py b/optuna_dashboard/_note.py index 95853a3d..507e3d8f 100644 --- a/optuna_dashboard/_note.py +++ b/optuna_dashboard/_note.py @@ -110,6 +110,19 @@ def note_str_key_prefix(trial_id: Optional[int]) -> str: return f"dashboard:{trial_id}:note_str:" +def copy_notes(storage: BaseStorage, src_study: optuna.Study, dst_study: optuna.Study) -> None: + system_attrs = storage.get_study_system_attrs(study_id=src_study._study_id) + + # Copy individual trial notes + for src_trial, dst_trial in zip(src_study.get_trials(), dst_study.get_trials()): + note = get_note_from_system_attrs(system_attrs, src_trial._trial_id)["body"] + save_note_with_version(storage, dst_study._study_id, dst_trial._trial_id, 0, note) + + # Copy study note + note = get_note_from_system_attrs(system_attrs, None)["body"] + save_note_with_version(storage, dst_study._study_id, None, 0, note) + + def get_note_from_system_attrs(system_attrs: dict[str, Any], trial_id: Optional[int]) -> NoteType: if note_ver_key(trial_id) not in system_attrs: return { diff --git a/python_tests/test_note.py b/python_tests/test_note.py index 7cc02cd4..592cf3e7 100644 --- a/python_tests/test_note.py +++ b/python_tests/test_note.py @@ -53,3 +53,25 @@ class NoteTestCase(TestCase): note_dict = note.get_note_from_system_attrs(system_attrs, trial._trial_id) self.assertEqual(note_dict["body"], body) self.assertEqual(note_dict["version"], expected_ver) + + def test_copy_notes(self) -> None: + old_study = optuna.create_study() + old_trials = [ + old_study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) for _ in range(2) + ] + storage = old_study._storage + + notes = ["trial 0", "trial 1"] + for trial, body in zip(old_trials, notes): + save_note(trial, body) + save_note(old_study, "Study") + + new_study = optuna.create_study(storage=storage, directions=old_study.directions) + new_study.add_trials(old_study.get_trials(deepcopy=False)) + + note.copy_notes(storage, old_study, new_study) + system_attrs = new_study._storage.get_study_system_attrs(new_study._study_id) + for new_trial, body in zip(new_study.get_trials(), notes): + actual = note.get_note_from_system_attrs(system_attrs, new_trial._trial_id) + self.assertEqual(actual["body"], body) + self.assertEqual(get_note(new_study), "Study")