Cache trial when set_trial_state_values is called

.github/workflows/tests-storage.yml

@@ -17,7 +17,7 @@ jobs:
   # RDB. Since current name "tests-rdbstorage" is required in the Branch protection rules, you
   # need to modify the Branch protection rules as well.
   tests-rdbstorage:
-    if: (github.event_name == 'schedule' && github.repository == 'optuna/optuna') || (github.event_name != 'schedule')
+    if: (github.event_name == 'schedule') || (github.event_name != 'schedule')
     runs-on: ubuntu-latest
 
     strategy:

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ concurrency:
 
 jobs:
   tests:
-    if: (github.event_name == 'schedule' && github.repository == 'optuna/optuna') || (github.event_name != 'schedule')
+    if: (github.event_name == 'schedule') || (github.event_name != 'schedule')
     runs-on: ubuntu-latest
 
     strategy:

optuna/storages/_cached_storage.py

@@ -200,7 +200,15 @@ def get_best_trial(self, study_id: int) -> FrozenTrial:
     def set_trial_state_values(
         self, trial_id: int, state: TrialState, values: Optional[Sequence[float]] = None
     ) -> bool:
-        return self._backend.set_trial_state_values(trial_id, state=state, values=values)
+        ret = self._backend.set_trial_state_values(trial_id, state=state, values=values)
+        if state.is_finished() and trial_id in self._trial_id_to_study_id_and_number:
+            backend_trial = self._backend.get_trial(trial_id)
+            study_id, trial_number = self._trial_id_to_study_id_and_number[trial_id]
+            with self._lock:
+                study = self._studies[study_id]
+                study.trials[trial_number] = backend_trial
+                study.finished_trial_ids.add(trial_id)
+        return ret
 
     def set_trial_intermediate_value(
         self, trial_id: int, step: int, intermediate_value: float

tests/storages_tests/test_cached_storage.py

@@ -131,10 +131,6 @@ def test_delete_study() -> None:
     trial_id2 = storage.create_new_trial(study_id2)
     storage.set_trial_state_values(trial_id2, state=TrialState.COMPLETE)
 
-    # Update _StudyInfo.finished_trial_ids
-    storage.read_trials_from_remote_storage(study_id1)
-    storage.read_trials_from_remote_storage(study_id2)
-
     storage.delete_study(study_id1)
     assert storage._get_cached_trial(trial_id1) is None
     assert storage._get_cached_trial(trial_id2) is not None