diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 9ce820b7..32e2a431 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -164,7 +164,7 @@ def create_app( if new_study_summary is None: response.status = 500 return {"reason": "Failed to load the new study"} - + note.transfer_notes(storage, src_study, dst_study) storage.delete_study(src_study._study_id) response.status = 201 diff --git a/optuna_dashboard/_note.py b/optuna_dashboard/_note.py index e510ddfd..6107d56c 100644 --- a/optuna_dashboard/_note.py +++ b/optuna_dashboard/_note.py @@ -109,6 +109,7 @@ def note_str_key_prefix(trial_id: Optional[int]) -> str: return prefix return f"dashboard:{trial_id}:note_str:" + def transfer_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) @@ -120,11 +121,12 @@ def transfer_notes(storage: BaseStorage, src_study: optuna.Study, dst_study: opt # Transfer individual trial notes for src_trial, dst_trial in zip(src_study.get_trials(), dst_study.get_trials()): transfer(src_trial._trial_id, dst_trial._trial_id) - + # Transfer study note NO_SRC_TRIAL, NO_DST_TRIAL = None, None transfer(NO_SRC_TRIAL, NO_DST_TRIAL) + 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 { @@ -146,7 +148,10 @@ def version_is_incremented( db_note_ver = system_attrs.get(note_ver_key(trial_id), 0) return req_note_ver == db_note_ver + 1 -def all_trial_notes(storage: BaseStorage, study_id: int, trial_id: Optional[int]) -> dict[str, str]: + +def all_trial_notes( + storage: BaseStorage, study_id: int, trial_id: Optional[int] +) -> dict[str, str]: all_note_attrs: dict[str, str] = { key: value for key, value in storage.get_study_system_attrs(study_id).items() @@ -154,6 +159,7 @@ def all_trial_notes(storage: BaseStorage, study_id: int, trial_id: Optional[int] } return all_note_attrs + def save_note_with_version( storage: BaseStorage, study_id: int, trial_id: Optional[int], ver: int, body: str ) -> None: @@ -169,19 +175,21 @@ def save_note_with_version( for i in range(len(attrs), len(all_note_attrs)): storage.set_study_system_attr(study_id, f"{note_str_key_prefix(trial_id)}{i}", "") + def delete_study_notes(storage: BaseStorage, study_id): - study = storage.get_study_name_from_id(study_id) for trial in storage.get_all_trials(study_id): delete_notes(storage, study_id, trial._trial_id) - + delete_notes(storage, study_id, None) + def delete_notes(storage: BaseStorage, study_id: int, trial_id: Optional[int]) -> None: all_note_attrs = all_trial_notes(storage, study_id, trial_id) for i in range(len(all_note_attrs)): storage.set_study_system_attr(study_id, f"{note_str_key_prefix(trial_id)}{i}", "") + def split_body(note_str: str, trial_id: Optional[int]) -> dict[str, str]: note_len = len(note_str) attrs = {} diff --git a/python_tests/test_note.py b/python_tests/test_note.py index 4997d16a..55c8f610 100644 --- a/python_tests/test_note.py +++ b/python_tests/test_note.py @@ -56,41 +56,70 @@ class NoteTestCase(TestCase): def test_delete_notes_trial(self) -> None: study = optuna.create_study() - trial_1 = study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) - trial_2 = study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) + trials = [ + study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) for _ in range(2) + ] storage = study._storage - for trial, body in [(trial_1, "version 1"), (trial_2, "version 2")]: - save_note(trial, body) + notes = ["trial 0", "trial 1"] + for trial, body in zip(trials, notes): + with self.subTest(body): + save_note(trial, body) - # first assert existence - actual = get_note(trial) - self.assertEqual(actual, body) + self.assertEqual(get_note(trial), body) - # delete - note.delete_notes(storage, study._study_id, trial._trial_id) + note.delete_notes(storage, study._study_id, trial._trial_id) - # assert deletion - actual = get_note(trial) - self.assertEqual(actual, "") + self.assertEqual(get_note(trial), "") + def test_delete_notes_study(self) -> None: + study = optuna.create_study() + trials = [ + study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) for _ in range(2) + ] + storage = study._storage + + notes = ["trial 0", "trial 1"] + for trial, body in zip(trials, notes): + with self.subTest(body): + save_note(trial, body) + + actual = get_note(trial) + self.assertEqual(actual, body) + + save_note(study, "Study note") + actual = get_note(study) + self.assertEqual(actual, "Study note") note.delete_study_notes(storage, study._study_id) - def test_delete_notes_study(self) -> None: - pass + for trial in trials: + self.assertEqual(get_note(trial), "") + self.assertEqual(get_note(study), "") def test_transfer_notes(self) -> None: - study = optuna.create_study() - trial_1 = study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) - trial_2 = study.ask({"x2": optuna.distributions.FloatDistribution(0, 10)}) - storage = study._storage + old_study = optuna.create_study() + old_trials = [ + old_study.ask({"x1": optuna.distributions.FloatDistribution(0, 10)}) for _ in range(2) + ] + storage = old_study._storage - save_note(trial_1, "trial 1") - save_note(trial_2, "trial 2") + 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=study.directions - ) - note.transfer_notes(storage, study, new_study) + new_study = optuna.create_study(storage=storage, directions=old_study.directions) + new_study.add_trials(old_study.get_trials(deepcopy=False)) + note.transfer_notes(storage, old_study, new_study) + + for old_trial in old_trials: + self.assertEqual(get_note(old_trial), "") + self.assertEqual(get_note(old_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")