We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optuna/storages/_heartbeat.py`
Content:
```
1 import abc
2 import copy
3 from threading import Event
4 from threading import Thread
5 from types import TracebackType
6 from typing import Callable
7 from typing import List
8 from typing import Optional
9 from typing import Type
10 
11 import optuna
12 from optuna._experimental import experimental_func
13 from optuna.storages import BaseStorage
14 from optuna.trial import FrozenTrial
15 from optuna.trial import TrialState
16 
17 
18 class BaseHeartbeat(metaclass=abc.ABCMeta):
19     """Base class for heartbeat.
20 
21     This class is not supposed to be directly accessed by library users.
22 
23     The heartbeat mechanism periodically checks whether each trial process is alive during an
24     optimization loop. To support this mechanism, the methods of
25     :class:`~optuna.storages._heartbeat.BaseHeartbeat` is implemented for the target database
26     backend, typically with multiple inheritance of :class:`~optuna.storages._base.BaseStorage`
27     and :class:`~optuna.storages._heartbeat.BaseHeartbeat`.
28 
29     .. seealso::
30         See :class:`~optuna.storages.RDBStorage`, where the backend supports heartbeat.
31     """
32 
33     @abc.abstractmethod
34     def record_heartbeat(self, trial_id: int) -> None:
35         """Record the heartbeat of the trial.
36 
37         Args:
38             trial_id:
39                 ID of the trial.
40         """
41         raise NotImplementedError()
42 
43     @abc.abstractmethod
44     def _get_stale_trial_ids(self, study_id: int) -> List[int]:
45         """Get the stale trial ids of the study.
46 
47         Args:
48             study_id:
49                 ID of the study.
50         Returns:
51             List of IDs of trials whose heartbeat has not been updated for a long time.
52         """
53         raise NotImplementedError()
54 
55     @abc.abstractmethod
56     def get_heartbeat_interval(self) -> Optional[int]:
57         """Get the heartbeat interval if it is set.
58 
59         Returns:
60             The heartbeat interval if it is set, otherwise :obj:`None`.
61         """
62         raise NotImplementedError()
63 
64     @abc.abstractmethod
65     def get_failed_trial_callback(self) -> Optional[Callable[["optuna.Study", FrozenTrial], None]]:
66         """Get the failed trial callback function.
67 
68         Returns:
69             The failed trial callback function if it is set, otherwise :obj:`None`.
70         """
71         raise NotImplementedError()
72 
73 
74 class BaseHeartbeatThread(metaclass=abc.ABCMeta):
75     def __enter__(self) -> None:
76         self.start()
77 
78     def __exit__(
79         self,
80         exc_type: Optional[Type[Exception]],
81         exc_value: Optional[Exception],
82         traceback: Optional[TracebackType],
83     ) -> None:
84         self.join()
85 
86     @abc.abstractmethod
87     def start(self) -> None:
88         raise NotImplementedError()
89 
90     @abc.abstractmethod
91     def join(self) -> None:
92         raise NotImplementedError()
93 
94 
95 class NullHeartbeatThread(BaseHeartbeatThread):
96     def __init__(self) -> None:
97         pass
98 
99     def start(self) -> None:
100         pass
101 
102     def join(self) -> None:
103         pass
104 
105 
106 class HeartbeatThread(BaseHeartbeatThread):
107     def __init__(self, trial_id: int, heartbeat: BaseHeartbeat) -> None:
108         self._trial_id = trial_id
109         self._heartbeat = heartbeat
110         self._thread: Optional[Thread] = None
111         self._stop_event: Optional[Event] = None
112 
113     def start(self) -> None:
114         self._stop_event = Event()
115         self._thread = Thread(
116             target=self._record_heartbeat, args=(self._trial_id, self._heartbeat, self._stop_event)
117         )
118         self._thread.start()
119 
120     def join(self) -> None:
121         assert self._stop_event is not None
122         assert self._thread is not None
123         self._stop_event.set()
124         self._thread.join()
125 
126     @staticmethod
127     def _record_heartbeat(trial_id: int, heartbeat: BaseHeartbeat, stop_event: Event) -> None:
128         heartbeat_interval = heartbeat.get_heartbeat_interval()
129         assert heartbeat_interval is not None
130         while True:
131             heartbeat.record_heartbeat(trial_id)
132             if stop_event.wait(timeout=heartbeat_interval):
133                 return
134 
135 
136 def get_heartbeat_thread(trial_id: int, storage: BaseStorage) -> BaseHeartbeatThread:
137     if is_heartbeat_enabled(storage):
138         assert isinstance(storage, BaseHeartbeat)
139         return HeartbeatThread(trial_id, storage)
140     else:
141         return NullHeartbeatThread()
142 
143 
144 @experimental_func("2.9.0")
145 def fail_stale_trials(study: "optuna.Study") -> None:
146     """Fail stale trials and run their failure callbacks.
147 
148     The running trials whose heartbeat has not been updated for a long time will be failed,
149     that is, those states will be changed to :obj:`~optuna.trial.TrialState.FAIL`.
150 
151     .. seealso::
152 
153         See :class:`~optuna.storages.RDBStorage`.
154 
155     Args:
156         study:
157             Study holding the trials to check.
158     """
159     storage = study._storage
160 
161     if not isinstance(storage, BaseHeartbeat):
162         return
163 
164     if not is_heartbeat_enabled(storage):
165         return
166 
167     failed_trial_ids = []
168     for trial_id in storage._get_stale_trial_ids(study._study_id):
169         if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):
170             failed_trial_ids.append(trial_id)
171 
172     failed_trial_callback = storage.get_failed_trial_callback()
173     if failed_trial_callback is not None:
174         for trial_id in failed_trial_ids:
175             failed_trial = copy.deepcopy(storage.get_trial(trial_id))
176             failed_trial_callback(study, failed_trial)
177 
178 
179 def is_heartbeat_enabled(storage: BaseStorage) -> bool:
180     """Check whether the storage enables the heartbeat.
181 
182     Returns:
183         :obj:`True` if the storage also inherits :class:`~optuna.storages._heartbeat.BaseHeartbeat`
184         and the return value of :meth:`~optuna.storages.BaseStorage.get_heartbeat_interval` is an
185         integer, otherwise :obj:`False`.
186     """
187     return isinstance(storage, BaseHeartbeat) and storage.get_heartbeat_interval() is not None
188
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
 
     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:

{"golden_diff": "diff --git a/optuna/storages/_heartbeat.py b/optuna/storages/_heartbeat.py\n--- a/optuna/storages/_heartbeat.py\n+++ b/optuna/storages/_heartbeat.py\n@@ -166,8 +166,12 @@\n \n     failed_trial_ids = []\n     for trial_id in storage._get_stale_trial_ids(study._study_id):\n-        if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):\n-            failed_trial_ids.append(trial_id)\n+        try:\n+            if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):\n+                failed_trial_ids.append(trial_id)\n+        except RuntimeError:\n+            # If another process fails the trial, the storage raises RuntimeError.\n+            pass\n \n     failed_trial_callback = storage.get_failed_trial_callback()\n     if failed_trial_callback is not None:\n", "issue": "RuntimeError caused by fail processing after liveness judgment in distributed optimization\n### Expected behavior\r\n\r\nOptuna 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.\r\n\r\nThis issue aims to discuss and address a solution to this problem.\r\n\r\n### Environment\r\n\r\n- Optuna version: 3.2.0\r\n- Python version: 3.9\r\n- OS: Ubuntu 22.04\r\n- (Optional) Other libraries and their versions:\r\n\r\n\r\n### Error messages, stack traces, or logs\r\n\r\n```shell\r\nTraceback (most recent call last):\r\n  File \"<string>\", line 3, in <module>\r\n  File \"xxxxx\", line 149, in chiran_study_search_task\r\n    study.search(n_trials)\r\n  File \"xxxxx\", line 170, in search\r\n    self._study.optimize(self._objective, callbacks=[MaxTrialsCallback(n_trials)])\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/study/study.py\", line 443, in optimize\r\n    _optimize(\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py\", line 66, in _optimize\r\n    _optimize_sequential(\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py\", line 163, in _optimize_sequential\r\n    frozen_trial = _run_trial(study, func, catch)\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/study/_optimize.py\", line 189, in _run_trial\r\n    optuna.storages.fail_stale_trials(study)\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/_experimental.py\", line 74, in wrapper\r\n    return func(*args, **kwargs)\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/storages/_heartbeat.py\", line 169, in fail_stale_trials\r\n    if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/storages/_cached_storage.py\", line 187, in set_trial_state_values\r\n    return self._backend.set_trial_state_values(trial_id, state=state, values=values)\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/storages/_rdb/storage.py\", line 640, in set_trial_state_values\r\n    self.check_trial_is_updatable(trial_id, trial.state)\r\n  File \"/usr/local/lib/python3.9/dist-packages/optuna/storages/_base.py\", line 626, in check_trial_is_updatable\r\n    raise RuntimeError(\r\nRuntimeError: Trial#63444 has already finished and can not be updated.\r\n```\r\n\r\n\r\n### Steps to reproduce\r\n\r\nRun Optuna optimization with 50 parallelism in k8s cluster environments, in which the pods are preempted frequently.\r\n\r\n### Additional context (optional)\r\n\r\nAs 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.\r\n\r\nAnother 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.\n", "before_files": [{"content": "import abc\nimport copy\nfrom threading import Event\nfrom threading import Thread\nfrom types import TracebackType\nfrom typing import Callable\nfrom typing import List\nfrom typing import Optional\nfrom typing import Type\n\nimport optuna\nfrom optuna._experimental import experimental_func\nfrom optuna.storages import BaseStorage\nfrom optuna.trial import FrozenTrial\nfrom optuna.trial import TrialState\n\n\nclass BaseHeartbeat(metaclass=abc.ABCMeta):\n    \"\"\"Base class for heartbeat.\n\n    This class is not supposed to be directly accessed by library users.\n\n    The heartbeat mechanism periodically checks whether each trial process is alive during an\n    optimization loop. To support this mechanism, the methods of\n    :class:`~optuna.storages._heartbeat.BaseHeartbeat` is implemented for the target database\n    backend, typically with multiple inheritance of :class:`~optuna.storages._base.BaseStorage`\n    and :class:`~optuna.storages._heartbeat.BaseHeartbeat`.\n\n    .. seealso::\n        See :class:`~optuna.storages.RDBStorage`, where the backend supports heartbeat.\n    \"\"\"\n\n    @abc.abstractmethod\n    def record_heartbeat(self, trial_id: int) -> None:\n        \"\"\"Record the heartbeat of the trial.\n\n        Args:\n            trial_id:\n                ID of the trial.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def _get_stale_trial_ids(self, study_id: int) -> List[int]:\n        \"\"\"Get the stale trial ids of the study.\n\n        Args:\n            study_id:\n                ID of the study.\n        Returns:\n            List of IDs of trials whose heartbeat has not been updated for a long time.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def get_heartbeat_interval(self) -> Optional[int]:\n        \"\"\"Get the heartbeat interval if it is set.\n\n        Returns:\n            The heartbeat interval if it is set, otherwise :obj:`None`.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def get_failed_trial_callback(self) -> Optional[Callable[[\"optuna.Study\", FrozenTrial], None]]:\n        \"\"\"Get the failed trial callback function.\n\n        Returns:\n            The failed trial callback function if it is set, otherwise :obj:`None`.\n        \"\"\"\n        raise NotImplementedError()\n\n\nclass BaseHeartbeatThread(metaclass=abc.ABCMeta):\n    def __enter__(self) -> None:\n        self.start()\n\n    def __exit__(\n        self,\n        exc_type: Optional[Type[Exception]],\n        exc_value: Optional[Exception],\n        traceback: Optional[TracebackType],\n    ) -> None:\n        self.join()\n\n    @abc.abstractmethod\n    def start(self) -> None:\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def join(self) -> None:\n        raise NotImplementedError()\n\n\nclass NullHeartbeatThread(BaseHeartbeatThread):\n    def __init__(self) -> None:\n        pass\n\n    def start(self) -> None:\n        pass\n\n    def join(self) -> None:\n        pass\n\n\nclass HeartbeatThread(BaseHeartbeatThread):\n    def __init__(self, trial_id: int, heartbeat: BaseHeartbeat) -> None:\n        self._trial_id = trial_id\n        self._heartbeat = heartbeat\n        self._thread: Optional[Thread] = None\n        self._stop_event: Optional[Event] = None\n\n    def start(self) -> None:\n        self._stop_event = Event()\n        self._thread = Thread(\n            target=self._record_heartbeat, args=(self._trial_id, self._heartbeat, self._stop_event)\n        )\n        self._thread.start()\n\n    def join(self) -> None:\n        assert self._stop_event is not None\n        assert self._thread is not None\n        self._stop_event.set()\n        self._thread.join()\n\n    @staticmethod\n    def _record_heartbeat(trial_id: int, heartbeat: BaseHeartbeat, stop_event: Event) -> None:\n        heartbeat_interval = heartbeat.get_heartbeat_interval()\n        assert heartbeat_interval is not None\n        while True:\n            heartbeat.record_heartbeat(trial_id)\n            if stop_event.wait(timeout=heartbeat_interval):\n                return\n\n\ndef get_heartbeat_thread(trial_id: int, storage: BaseStorage) -> BaseHeartbeatThread:\n    if is_heartbeat_enabled(storage):\n        assert isinstance(storage, BaseHeartbeat)\n        return HeartbeatThread(trial_id, storage)\n    else:\n        return NullHeartbeatThread()\n\n\n@experimental_func(\"2.9.0\")\ndef fail_stale_trials(study: \"optuna.Study\") -> None:\n    \"\"\"Fail stale trials and run their failure callbacks.\n\n    The running trials whose heartbeat has not been updated for a long time will be failed,\n    that is, those states will be changed to :obj:`~optuna.trial.TrialState.FAIL`.\n\n    .. seealso::\n\n        See :class:`~optuna.storages.RDBStorage`.\n\n    Args:\n        study:\n            Study holding the trials to check.\n    \"\"\"\n    storage = study._storage\n\n    if not isinstance(storage, BaseHeartbeat):\n        return\n\n    if not is_heartbeat_enabled(storage):\n        return\n\n    failed_trial_ids = []\n    for trial_id in storage._get_stale_trial_ids(study._study_id):\n        if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):\n            failed_trial_ids.append(trial_id)\n\n    failed_trial_callback = storage.get_failed_trial_callback()\n    if failed_trial_callback is not None:\n        for trial_id in failed_trial_ids:\n            failed_trial = copy.deepcopy(storage.get_trial(trial_id))\n            failed_trial_callback(study, failed_trial)\n\n\ndef is_heartbeat_enabled(storage: BaseStorage) -> bool:\n    \"\"\"Check whether the storage enables the heartbeat.\n\n    Returns:\n        :obj:`True` if the storage also inherits :class:`~optuna.storages._heartbeat.BaseHeartbeat`\n        and the return value of :meth:`~optuna.storages.BaseStorage.get_heartbeat_interval` is an\n        integer, otherwise :obj:`False`.\n    \"\"\"\n    return isinstance(storage, BaseHeartbeat) and storage.get_heartbeat_interval() is not None\n", "path": "optuna/storages/_heartbeat.py"}], "after_files": [{"content": "import abc\nimport copy\nfrom threading import Event\nfrom threading import Thread\nfrom types import TracebackType\nfrom typing import Callable\nfrom typing import List\nfrom typing import Optional\nfrom typing import Type\n\nimport optuna\nfrom optuna._experimental import experimental_func\nfrom optuna.storages import BaseStorage\nfrom optuna.trial import FrozenTrial\nfrom optuna.trial import TrialState\n\n\nclass BaseHeartbeat(metaclass=abc.ABCMeta):\n    \"\"\"Base class for heartbeat.\n\n    This class is not supposed to be directly accessed by library users.\n\n    The heartbeat mechanism periodically checks whether each trial process is alive during an\n    optimization loop. To support this mechanism, the methods of\n    :class:`~optuna.storages._heartbeat.BaseHeartbeat` is implemented for the target database\n    backend, typically with multiple inheritance of :class:`~optuna.storages._base.BaseStorage`\n    and :class:`~optuna.storages._heartbeat.BaseHeartbeat`.\n\n    .. seealso::\n        See :class:`~optuna.storages.RDBStorage`, where the backend supports heartbeat.\n    \"\"\"\n\n    @abc.abstractmethod\n    def record_heartbeat(self, trial_id: int) -> None:\n        \"\"\"Record the heartbeat of the trial.\n\n        Args:\n            trial_id:\n                ID of the trial.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def _get_stale_trial_ids(self, study_id: int) -> List[int]:\n        \"\"\"Get the stale trial ids of the study.\n\n        Args:\n            study_id:\n                ID of the study.\n        Returns:\n            List of IDs of trials whose heartbeat has not been updated for a long time.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def get_heartbeat_interval(self) -> Optional[int]:\n        \"\"\"Get the heartbeat interval if it is set.\n\n        Returns:\n            The heartbeat interval if it is set, otherwise :obj:`None`.\n        \"\"\"\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def get_failed_trial_callback(self) -> Optional[Callable[[\"optuna.Study\", FrozenTrial], None]]:\n        \"\"\"Get the failed trial callback function.\n\n        Returns:\n            The failed trial callback function if it is set, otherwise :obj:`None`.\n        \"\"\"\n        raise NotImplementedError()\n\n\nclass BaseHeartbeatThread(metaclass=abc.ABCMeta):\n    def __enter__(self) -> None:\n        self.start()\n\n    def __exit__(\n        self,\n        exc_type: Optional[Type[Exception]],\n        exc_value: Optional[Exception],\n        traceback: Optional[TracebackType],\n    ) -> None:\n        self.join()\n\n    @abc.abstractmethod\n    def start(self) -> None:\n        raise NotImplementedError()\n\n    @abc.abstractmethod\n    def join(self) -> None:\n        raise NotImplementedError()\n\n\nclass NullHeartbeatThread(BaseHeartbeatThread):\n    def __init__(self) -> None:\n        pass\n\n    def start(self) -> None:\n        pass\n\n    def join(self) -> None:\n        pass\n\n\nclass HeartbeatThread(BaseHeartbeatThread):\n    def __init__(self, trial_id: int, heartbeat: BaseHeartbeat) -> None:\n        self._trial_id = trial_id\n        self._heartbeat = heartbeat\n        self._thread: Optional[Thread] = None\n        self._stop_event: Optional[Event] = None\n\n    def start(self) -> None:\n        self._stop_event = Event()\n        self._thread = Thread(\n            target=self._record_heartbeat, args=(self._trial_id, self._heartbeat, self._stop_event)\n        )\n        self._thread.start()\n\n    def join(self) -> None:\n        assert self._stop_event is not None\n        assert self._thread is not None\n        self._stop_event.set()\n        self._thread.join()\n\n    @staticmethod\n    def _record_heartbeat(trial_id: int, heartbeat: BaseHeartbeat, stop_event: Event) -> None:\n        heartbeat_interval = heartbeat.get_heartbeat_interval()\n        assert heartbeat_interval is not None\n        while True:\n            heartbeat.record_heartbeat(trial_id)\n            if stop_event.wait(timeout=heartbeat_interval):\n                return\n\n\ndef get_heartbeat_thread(trial_id: int, storage: BaseStorage) -> BaseHeartbeatThread:\n    if is_heartbeat_enabled(storage):\n        assert isinstance(storage, BaseHeartbeat)\n        return HeartbeatThread(trial_id, storage)\n    else:\n        return NullHeartbeatThread()\n\n\n@experimental_func(\"2.9.0\")\ndef fail_stale_trials(study: \"optuna.Study\") -> None:\n    \"\"\"Fail stale trials and run their failure callbacks.\n\n    The running trials whose heartbeat has not been updated for a long time will be failed,\n    that is, those states will be changed to :obj:`~optuna.trial.TrialState.FAIL`.\n\n    .. seealso::\n\n        See :class:`~optuna.storages.RDBStorage`.\n\n    Args:\n        study:\n            Study holding the trials to check.\n    \"\"\"\n    storage = study._storage\n\n    if not isinstance(storage, BaseHeartbeat):\n        return\n\n    if not is_heartbeat_enabled(storage):\n        return\n\n    failed_trial_ids = []\n    for trial_id in storage._get_stale_trial_ids(study._study_id):\n        try:\n            if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):\n                failed_trial_ids.append(trial_id)\n        except RuntimeError:\n            # If another process fails the trial, the storage raises RuntimeError.\n            pass\n\n    failed_trial_callback = storage.get_failed_trial_callback()\n    if failed_trial_callback is not None:\n        for trial_id in failed_trial_ids:\n            failed_trial = copy.deepcopy(storage.get_trial(trial_id))\n            failed_trial_callback(study, failed_trial)\n\n\ndef is_heartbeat_enabled(storage: BaseStorage) -> bool:\n    \"\"\"Check whether the storage enables the heartbeat.\n\n    Returns:\n        :obj:`True` if the storage also inherits :class:`~optuna.storages._heartbeat.BaseHeartbeat`\n        and the return value of :meth:`~optuna.storages.BaseStorage.get_heartbeat_interval` is an\n        integer, otherwise :obj:`False`.\n    \"\"\"\n    return isinstance(storage, BaseHeartbeat) and storage.get_heartbeat_interval() is not None\n", "path": "optuna/storages/_heartbeat.py"}]}