Repository: optuna/optuna

Issue:
RuntimeError caused by fail processing after liveness judgment in distributed optimization
### Expected behavior

Optuna incorporates a Heartbeat mechanism, which evaluates trials by detecting those that haven't recorded a heartbeat within a certain time period and transitions them from "Running" to "Fail". However, the function responsible for this liveness judgment, `_get_stale_trial_ids` within `BaseHeartbeat.fail_stale_trials`, is not atomic. Consequently, during distributed optimization, different processes might acquire the same trial ID, leading to parallel invocations of `storage.set_trial_state_values` for the same trial ID. This can result in a `RuntimeError`, as Optuna triggers a `RuntimeError` when calling `storage.set_trial_state_values` for trials that have already completed.

This issue aims to discuss and address a solution to this problem.

### Environment

- Optuna version: 3.2.0
- Python version: 3.9
- OS: Ubuntu 22.04
- (Optional) Other libraries and their versions:


### Error messages, stack traces, or logs

```shell
Traceback (most recent call last):
  File "<string>", line 3, in <module>
  File "xxxxx", line 149, in chiran_study_search_task
    study.search(n_trials)
  File "xxxxx", line 170, in search
    self._study.optimize(self._objective, callbacks=[MaxTrialsCallback(n_trials)])
  File "/usr/local/lib/python3.9/dist-packages/optuna/study/study.py", line 443, in optimize
    _optimize(
  File "/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py", line 66, in _optimize
    _optimize_sequential(
  File "/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py", line 163, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
  File "/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py", line 189, in _run_trial
    optuna.storages.fail_stale_trials(study)
  File "/usr/local/lib/python3.9/dist-packages/optuna/_experimental.py", line 74, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/optuna/storages/_heartbeat.py", line 169, in fail_stale_trials
    if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):
  File "/usr/local/lib/python3.9/dist-packages/optuna/storages/_cached_storage.py", line 187, in set_trial_state_values
    return self._backend.set_trial_state_values(trial_id, state=state, values=values)
  File "/usr/local/lib/python3.9/dist-packages/optuna/storages/_rdb/storage.py", line 640, in set_trial_state_values
    self.check_trial_is_updatable(trial_id, trial.state)
  File "/usr/local/lib/python3.9/dist-packages/optuna/storages/_base.py", line 626, in check_trial_is_updatable
    raise RuntimeError(
RuntimeError: Trial#63444 has already finished and can not be updated.
```


### Steps to reproduce

Run Optuna optimization with 50 parallelism in k8s cluster environments, in which the pods are preempted frequently.

### Additional context (optional)

As a solution, one possible approach is to make the `_get_stale_trial_ids` call atomic for a single DB. However, this would involve introducing reference locks on the DB, potentially becoming a bottleneck as the parallelism of distributed optimization increases. Therefore, personally, I don't consider this to be an ideal solution.

Another solution could involve reducing the frequency of `fail_stale_trials` calls, rather than invoking it for every trial. This entails adding an argument like `trial_interval_of_heartbeat` to the storage that inherits from `BaseHeartbeat`, and invoking `fail_stale_trials` for trials at this interval. With settings where the level of parallelism is smaller than `trial_interval_of_heartbeat`, it is expected that the issue at hand would not occur.


Discussion:

Resolving patch:
diff --git a/optuna/storages/_heartbeat.py b/optuna/storages/_heartbeat.py
--- a/optuna/storages/_heartbeat.py
+++ b/optuna/storages/_heartbeat.py
@@ -166,8 +166,12 @@ def fail_stale_trials(study: "optuna.Study") -> None:
 
     failed_trial_ids = []
     for trial_id in storage._get_stale_trial_ids(study._study_id):
-        if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):
-            failed_trial_ids.append(trial_id)
+        try:
+            if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):
+                failed_trial_ids.append(trial_id)
+        except RuntimeError:
+            # If another process fails the trial, the storage raises RuntimeError.
+            pass
 
     failed_trial_callback = storage.get_failed_trial_callback()
     if failed_trial_callback is not None:


Test patch:
diff --git a/tests/storages_tests/test_heartbeat.py b/tests/storages_tests/test_heartbeat.py
--- a/tests/storages_tests/test_heartbeat.py
+++ b/tests/storages_tests/test_heartbeat.py
@@ -1,5 +1,7 @@
 import itertools
+import multiprocessing
 import time
+from typing import Any
 from typing import Optional
 from unittest.mock import Mock
 from unittest.mock import patch
@@ -235,6 +237,50 @@ def check_keep_trial_state_in_running(study: "optuna.Study") -> None:
         check_change_trial_state_to_fail(study)
 
 
+def run_fail_stale_trials(storage_url: str, sleep_time: int) -> None:
+    heartbeat_interval = 1
+    grace_period = 2
+
+    storage = RDBStorage(storage_url)
+    storage.heartbeat_interval = heartbeat_interval
+    storage.grace_period = grace_period
+
+    original_set_trial_state_values = storage.set_trial_state_values
+
+    def _set_trial_state_values(*args: Any, **kwargs: Any) -> bool:
+        # The second process fails to set state due to the race condition.
+        time.sleep(sleep_time)
+        return original_set_trial_state_values(*args, **kwargs)
+
+    storage.set_trial_state_values = _set_trial_state_values  # type: ignore[method-assign]
+
+    study = optuna.load_study(study_name=None, storage=storage)
+    optuna.storages.fail_stale_trials(study)
+
+
+def test_fail_stale_trials_with_race_condition() -> None:
+    grace_period = 2
+    storage_mode = "sqlite"
+
+    with StorageSupplier(storage_mode) as storage:
+        assert isinstance(storage, RDBStorage)
+
+        study = optuna.create_study(storage=storage)
+
+        trial = study.ask()
+        storage.record_heartbeat(trial._trial_id)
+        time.sleep(grace_period + 1)
+        p1 = multiprocessing.Process(target=run_fail_stale_trials, args=(storage.url, 1))
+        p1.start()
+        p2 = multiprocessing.Process(target=run_fail_stale_trials, args=(storage.url, 2))
+        p2.start()
+        p1.join()
+        p2.join()
+        assert p1.exitcode == 0
+        assert p2.exitcode == 0
+        assert study.trials[0].state is TrialState.FAIL
+
+
 def test_get_stale_trial_ids() -> None:
     storage_mode = "sqlite"
     heartbeat_interval = 1