Repository: optuna/optuna
Pull request: Redis Heartbeat
<!-- Thank you for creating a pull request! In general, we merge your pull requests after they get two or more approvals. -->

## Motivation
Fix #2773. Sorry for not waiting for your replies there, but I could not wait to start coding :smile: 

## Description of the changes
This PR adds heartbeat support to the Redis storage backend and unifies some of the heartbeat related tests for RDB and Redis backends. Detailed list of changes:

- Heartbeats are saved in a hash (i.e. dict) per study and not per trial. While this requires an additional lookup of the study ID in `set_trial_state()` and `record_heartbeat()`, it significantly improves the performance of `fail_stale_trials()` for studies with many non-running trials.
- Heartbeats are removed when a trial's state is updated on completion. As empty hashes are deleted in Redis automatically, there should be no permanent storage usage after finishing all running trials. However, a user can call `record_heartbeat` out of order, but as he/she is not supposed to call backend methods anyway and the RDB storage has the same problem, I ignored this. After all, looking up the trial state in `record_heartbeat()` would cause a performance cost.
- I moved the tests from the RDB storage tests to the general storage tests and added a new list `STORAGE_MODES_HEARTBEAT`. Additionally, `StorageSupplier` now supports passing `kwargs` to the backends to unify the setting of `heartbeat_interval`, `grace_period` and `failed_trial_callback`. In this process, I lost one single test which is redundant in my opinion, i.e. checking the return value of `_get_stale_trials()` in `test_fail_stale_trials()`.


Problem statement:
Redis Heartbeart
Add heartbeat support to the Redis storage backend similar to the RDB storage.

## Motivation

While the RDB storage backend supports the heartbeat functionality, the Redis storage does not. This prevents the Redis backend from being a drop-in replacement for the RDB backend.

## Description

Implement the methods of `BaseStorage` in `RedisStorage`.

## Alternatives

Move the heartbeat functionality out of the individual storage classes into `BaseStorage`, e.g. via trial system attributes. Although this provides some simplifications, I would not recommend this because deleting stale trials can be implemented more efficiently by the backend.

## Additional context

I would like to contribute this feature, if you think it fits into the library.

Hints from issue discussion:

Commit message: Simplify test

Patch to review:
diff --git a/optuna/storages/_redis.py b/optuna/storages/_redis.py
index 31abb0f8285..e7b95c27863 100644
--- a/optuna/storages/_redis.py
+++ b/optuna/storages/_redis.py
@@ -1,7 +1,9 @@
 import copy
 from datetime import datetime
 import pickle
+import time
 from typing import Any
+from typing import Callable
 from typing import cast
 from typing import Dict
 from typing import List
@@ -59,19 +61,50 @@ def objective(trial):
 
     Args:
         url: URL of the redis storage, password and db are optional. (ie: redis://localhost:6379)
+        heartbeat_interval:
+            Interval to record the heartbeat. It is recorded every ``interval`` seconds.
+        grace_period:
+            Grace period before a running trial is failed from the last heartbeat.
+            If it is :obj:`None`, the grace period will be `2 * heartbeat_interval`.
+        failed_trial_callback:
+            A callback function that is invoked after failing each stale trial.
+            The function must accept two parameters with the following types in this order:
+            :class:`~optuna.study.Study` and :class:`~optuna.FrozenTrial`.
+
+            .. note::
+                The procedure to fail existing stale trials is called just before asking the
+                study for a new trial.
 
     .. note::
         If you use plan to use Redis as a storage mechanism for optuna,
         make sure Redis in installed and running.
         Please execute ``$ pip install -U redis`` to install redis python library.
+
+    Raises:
+        :exc:`ValueError`:
+            If the given `heartbeat_interval` or `grace_period` is not a positive integer.
     """
 
-    def __init__(self, url: str) -> None:
+    def __init__(
+        self,
+        url: str,
+        heartbeat_interval: Optional[int] = None,
+        grace_period: Optional[int] = None,
+        failed_trial_callback: Optional[Callable[["optuna.Study", FrozenTrial], None]] = None,
+    ) -> None:
 
         _imports.check()
 
+        if heartbeat_interval is not None and heartbeat_interval <= 0:
+            raise ValueError("The value of `heartbeat_interval` should be a positive integer.")
+        if grace_period is not None and grace_period <= 0:
+            raise ValueError("The value of `grace_period` should be a positive integer.")
+
         self._url = url
         self._redis = redis.Redis.from_url(url)
+        self._heartbeat_interval = heartbeat_interval
+        self._grace_period = grace_period
+        self._failed_trial_callback = failed_trial_callback
 
     def create_new_study(self, study_name: Optional[str] = None) -> int:
 
@@ -365,6 +398,9 @@ def set_trial_state(self, trial_id: int, state: TrialState) -> bool:
             trial.datetime_complete = datetime.now()
             self._redis.set(self._key_trial(trial_id), pickle.dumps(trial))
             self._update_cache(trial_id)
+
+            study_id = self.get_study_id_from_trial_id(trial_id)
+            self._redis.hdel(self._key_study_heartbeats(study_id), str(trial_id))
         else:
             self._redis.set(self._key_trial(trial_id), pickle.dumps(trial))
 
@@ -607,3 +643,48 @@ def _check_trial_id(self, trial_id: int) -> None:
 
         if not self._redis.exists(self._key_trial(trial_id)):
             raise KeyError("study_id {} does not exist.".format(trial_id))
+
+    def record_heartbeat(self, trial_id: int) -> None:
+        study_id = self.get_study_id_from_trial_id(trial_id)
+        self._redis.hset(
+            self._key_study_heartbeats(study_id), str(trial_id), pickle.dumps(time.time())
+        )
+
+    def fail_stale_trials(self, study_id: int) -> List[int]:
+        confirmed = []
+        for trial_id in self._get_stale_trial_ids(study_id):
+            if self.set_trial_state(trial_id, TrialState.FAIL):
+                confirmed.append(trial_id)
+
+        return confirmed
+
+    def _get_stale_trial_ids(self, study_id: int) -> List[int]:
+        assert self._heartbeat_interval is not None
+
+        if self._grace_period is None:
+            grace_period = 2 * self._heartbeat_interval
+        else:
+            grace_period = self._grace_period
+
+        current_time = time.time()
+        heartbeats = self._redis.hgetall(self._key_study_heartbeats(study_id))
+        stale = []
+        for trial_id_raw, last_heartbeat_raw in heartbeats.items():
+            last_heartbeat = pickle.loads(last_heartbeat_raw)
+            if current_time - last_heartbeat > grace_period:
+                trial_id = int(trial_id_raw)
+                stale.append(trial_id)
+        return stale
+
+    def _is_heartbeat_supported(self) -> bool:
+        return True
+
+    def get_heartbeat_interval(self) -> Optional[int]:
+        return self._heartbeat_interval
+
+    @staticmethod
+    def _key_study_heartbeats(study_id: int) -> str:
+        return "study_id:{:010d}:heartbeats".format(study_id)
+
+    def get_failed_trial_callback(self) -> Optional[Callable[["optuna.Study", FrozenTrial], None]]:
+        return self._failed_trial_callback
diff --git a/optuna/testing/storage.py b/optuna/testing/storage.py
index 04898b300af..24e25e4548c 100644
--- a/optuna/testing/storage.py
+++ b/optuna/testing/storage.py
@@ -17,35 +17,48 @@
     "redis",
 ]
 
+STORAGE_MODES_HEARTBEAT = [
+    "sqlite",
+    "cache",
+    "redis",
+]
+
 SQLITE3_TIMEOUT = 300
 
 
 class StorageSupplier(object):
-    def __init__(self, storage_specifier: str) -> None:
+    def __init__(self, storage_specifier: str, **kwargs: Any) -> None:
 
         self.storage_specifier = storage_specifier
         self.tempfile: Optional[IO[Any]] = None
+        self.extra_args = kwargs
 
     def __enter__(self) -> optuna.storages.BaseStorage:
 
         if self.storage_specifier == "inmemory":
+            if len(self.extra_args) > 0:
+                raise ValueError("InMemoryStorage does not accept any arguments!")
             return optuna.storages.InMemoryStorage()
         elif self.storage_specifier == "sqlite":
             self.tempfile = tempfile.NamedTemporaryFile()
             url = "sqlite:///{}".format(self.tempfile.name)
             return optuna.storages.RDBStorage(
-                url, engine_kwargs={"connect_args": {"timeout": SQLITE3_TIMEOUT}}
+                url,
+                engine_kwargs={"connect_args": {"timeout": SQLITE3_TIMEOUT}},
+                **self.extra_args,
             )
         elif self.storage_specifier == "cache":
             self.tempfile = tempfile.NamedTemporaryFile()
             url = "sqlite:///{}".format(self.tempfile.name)
             return optuna.storages._CachedStorage(
                 optuna.storages.RDBStorage(
-                    url, engine_kwargs={"connect_args": {"timeout": SQLITE3_TIMEOUT}}
+                    url,
+                    engine_kwargs={"connect_args": {"timeout": SQLITE3_TIMEOUT}},
+                    **self.extra_args,
                 )
             )
         elif self.storage_specifier == "redis":
-            storage = optuna.storages.RedisStorage("redis://localhost")
+            storage = optuna.storages.RedisStorage("redis://localhost", **self.extra_args)
             storage._redis = fakeredis.FakeStrictRedis()
             return storage
         else:
diff --git a/tests/storages_tests/rdb_tests/test_storage.py b/tests/storages_tests/rdb_tests/test_storage.py
index 6b9eb9be591..1b246bc7920 100644
--- a/tests/storages_tests/rdb_tests/test_storage.py
+++ b/tests/storages_tests/rdb_tests/test_storage.py
@@ -12,12 +12,10 @@
 from sqlalchemy.exc import IntegrityError
 
 from optuna import create_study
-from optuna import Study
 from optuna import version
 from optuna.distributions import CategoricalDistribution
 from optuna.distributions import UniformDistribution
 from optuna.storages import RDBStorage
-from optuna.storages import RetryFailedTrialCallback
 from optuna.storages._rdb.models import SCHEMA_VERSION
 from optuna.storages._rdb.models import TrialHeartbeatModel
 from optuna.storages._rdb.models import VersionInfoModel
@@ -332,101 +330,3 @@ def objective(trial: Trial) -> float:
         assert len(trial_heartbeats) == n_trials
         for i in range(n_trials - 1):
             assert (trial_heartbeats[i + 1] - trial_heartbeats[i]).seconds - sleep_sec <= 1
-
-
-def test_fail_stale_trials() -> None:
-
-    heartbeat_interval = 1
-    grace_period = 2
-
-    with StorageSupplier("sqlite") as storage:
-        assert isinstance(storage, RDBStorage)
-        storage.heartbeat_interval = heartbeat_interval
-        storage.grace_period = grace_period
-        study1 = create_study(storage=storage)
-        study2 = create_study(storage=storage)
-
-        trial1 = study1.ask()
-        trial2 = study2.ask()
-        storage.record_heartbeat(trial1._trial_id)
-        storage.record_heartbeat(trial2._trial_id)
-        time.sleep(grace_period + 1)
-
-        assert study1.trials[0].state is TrialState.RUNNING
-        assert study2.trials[0].state is TrialState.RUNNING
-
-        assert storage._get_stale_trial_ids(study1._study_id) == [study1.trials[0]._trial_id]
-
-        # Exceptions raised in spawned threads are caught by `_TestableThread`.
-        with patch("optuna.study._optimize.Thread", _TestableThread):
-            study1.optimize(lambda _: 1.0, n_trials=1)
-
-        assert study1.trials[0].state is TrialState.FAIL
-        assert study2.trials[0].state is TrialState.RUNNING
-
-
-def test_invalid_heartbeat_interval_and_grace_period() -> None:
-
-    with pytest.raises(ValueError):
-        _ = RDBStorage("sqlite:///:memory:", heartbeat_interval=-1)
-
-    with pytest.raises(ValueError):
-        _ = RDBStorage("sqlite:///:memory:", grace_period=-1)
-
-
-def test_failed_trial_callback() -> None:
-    heartbeat_interval = 1
-    grace_period = 2
-
-    def failed_trial_callback(study: Study, trial: FrozenTrial) -> None:
-        assert study.system_attrs["test"] == "A"
-        assert trial.system_attrs["test"] == "B"
-
-    with StorageSupplier("sqlite") as storage:
-        assert isinstance(storage, RDBStorage)
-        storage.heartbeat_interval = heartbeat_interval
-        storage.grace_period = grace_period
-        storage.failed_trial_callback = failed_trial_callback
-        study = create_study(storage=storage)
-        study.set_system_attr("test", "A")
-
-        trial = study.ask()
-        trial.set_system_attr("test", "B")
-        storage.record_heartbeat(trial._trial_id)
-        time.sleep(grace_period + 1)
-
-        # Exceptions raised in spawned threads are caught by `_TestableThread`.
-        with patch("optuna.study._optimize.Thread", _TestableThread):
-            with patch.object(storage, "failed_trial_callback", wraps=failed_trial_callback) as m:
-                study.optimize(lambda _: 1.0, n_trials=1)
-                m.assert_called_once()
-
-
-@pytest.mark.parametrize("max_retry", [None, 0, 1])
-def test_retry_failed_trial_callback(max_retry: Optional[int]) -> None:
-    heartbeat_interval = 1
-    grace_period = 2
-
-    with StorageSupplier("sqlite") as storage:
-        assert isinstance(storage, RDBStorage)
-        storage.heartbeat_interval = heartbeat_interval
-        storage.grace_period = grace_period
-
-        storage.failed_trial_callback = RetryFailedTrialCallback(max_retry=max_retry)
-        study = create_study(storage=storage)
-
-        trial = study.ask()
-        storage.record_heartbeat(trial._trial_id)
-        time.sleep(grace_period + 1)
-
-        # Exceptions raised in spawned threads are caught by `_TestableThread`.
-        with patch("optuna.study._optimize.Thread", _TestableThread):
-            study.optimize(lambda _: 1.0, n_trials=1)
-
-        # Test the last trial to see if it was a retry of the first trial or not.
-        # Test max_rety=None to see if trial is retried.
-        # Test max_rety=0 to see if no trials are retried.
-        # Test max_rety=1 to see if trial is retried.
-        assert RetryFailedTrialCallback.retried_trial_number(study.trials[1]) == (
-            None if max_retry == 0 else 0
-        )
diff --git a/tests/storages_tests/test_storages.py b/tests/storages_tests/test_storages.py
index 151b58aa747..0c36afe5cf6 100644
--- a/tests/storages_tests/test_storages.py
+++ b/tests/storages_tests/test_storages.py
@@ -1,17 +1,22 @@
 import copy
 from datetime import datetime
+import itertools
 import random
+import time
 from typing import Any
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
+from unittest.mock import Mock
 from unittest.mock import patch
 
 import pytest
 
 import optuna
+from optuna import Study
+from optuna._callbacks import RetryFailedTrialCallback
 from optuna.distributions import CategoricalDistribution
 from optuna.distributions import LogUniformDistribution
 from optuna.distributions import UniformDistribution
@@ -24,7 +29,9 @@
 from optuna.study._study_direction import StudyDirection
 from optuna.study._study_summary import StudySummary
 from optuna.testing.storage import STORAGE_MODES
+from optuna.testing.storage import STORAGE_MODES_HEARTBEAT
 from optuna.testing.storage import StorageSupplier
+from optuna.testing.threading import _TestableThread
 from optuna.trial import FrozenTrial
 from optuna.trial import TrialState
 
@@ -1087,3 +1094,113 @@ def test_get_trial_id_from_study_id_trial_number(storage_mode: str) -> None:
         assert trial_id == storage.get_trial_id_from_study_id_trial_number(
             study_id, trial_number=0
         )
+
+
+@pytest.mark.parametrize("storage_mode", STORAGE_MODES_HEARTBEAT)
+def test_fail_stale_trials(storage_mode: str) -> None:
+
+    heartbeat_interval = 1
+    grace_period = 2
+
+    with StorageSupplier(
+        storage_mode, heartbeat_interval=heartbeat_interval, grace_period=grace_period
+    ) as storage:
+        assert storage.is_heartbeat_enabled()
+
+        study1 = optuna.create_study(storage=storage)
+        study2 = optuna.create_study(storage=storage)
+
+        trial1 = study1.ask()
+        trial2 = study2.ask()
+        storage.record_heartbeat(trial1._trial_id)
+        storage.record_heartbeat(trial2._trial_id)
+        time.sleep(grace_period + 1)
+
+        assert study1.trials[0].state is TrialState.RUNNING
+        assert study2.trials[0].state is TrialState.RUNNING
+
+        # Exceptions raised in spawned threads are caught by `_TestableThread`.
+        with patch("optuna.study._optimize.Thread", _TestableThread):
+            study1.optimize(lambda _: 1.0, n_trials=1)
+
+        assert study1.trials[0].state is TrialState.FAIL
+        assert study2.trials[0].state is TrialState.RUNNING
+
+
+@pytest.mark.parametrize("storage_mode", STORAGE_MODES_HEARTBEAT)
+def test_invalid_heartbeat_interval_and_grace_period(storage_mode: str) -> None:
+
+    with pytest.raises(ValueError):
+        with StorageSupplier(storage_mode, heartbeat_interval=-1):
+            pass
+
+    with pytest.raises(ValueError):
+        with StorageSupplier(storage_mode, grace_period=-1):
+            pass
+
+
+@pytest.mark.parametrize("storage_mode", STORAGE_MODES_HEARTBEAT)
+def test_failed_trial_callback(storage_mode: str) -> None:
+    heartbeat_interval = 1
+    grace_period = 2
+
+    def _failed_trial_callback(study: Study, trial: FrozenTrial) -> None:
+        assert study.system_attrs["test"] == "A"
+        assert trial.system_attrs["test"] == "B"
+
+    failed_trial_callback = Mock(wraps=_failed_trial_callback)
+
+    with StorageSupplier(
+        storage_mode,
+        heartbeat_interval=heartbeat_interval,
+        grace_period=grace_period,
+        failed_trial_callback=failed_trial_callback,
+    ) as storage:
+        assert storage.is_heartbeat_enabled()
+
+        study = optuna.create_study(storage=storage)
+        study.set_system_attr("test", "A")
+
+        trial = study.ask()
+        trial.set_system_attr("test", "B")
+        storage.record_heartbeat(trial._trial_id)
+        time.sleep(grace_period + 1)
+
+        # Exceptions raised in spawned threads are caught by `_TestableThread`.
+        with patch("optuna.study._optimize.Thread", _TestableThread):
+            study.optimize(lambda _: 1.0, n_trials=1)
+            failed_trial_callback.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "storage_mode,max_retry", itertools.product(STORAGE_MODES_HEARTBEAT, [None, 0, 1])
+)
+def test_retry_failed_trial_callback(storage_mode: str, max_retry: Optional[int]) -> None:
+    heartbeat_interval = 1
+    grace_period = 2
+
+    with StorageSupplier(
+        storage_mode,
+        heartbeat_interval=heartbeat_interval,
+        grace_period=grace_period,
+        failed_trial_callback=RetryFailedTrialCallback(max_retry=max_retry),
+    ) as storage:
+        assert storage.is_heartbeat_enabled()
+
+        study = optuna.create_study(storage=storage)
+
+        trial = study.ask()
+        storage.record_heartbeat(trial._trial_id)
+        time.sleep(grace_period + 1)
+
+        # Exceptions raised in spawned threads are caught by `_TestableThread`.
+        with patch("optuna.study._optimize.Thread", _TestableThread):
+            study.optimize(lambda _: 1.0, n_trials=1)
+
+        # Test the last trial to see if it was a retry of the first trial or not.
+        # Test max_retry=None to see if trial is retried.
+        # Test max_retry=0 to see if no trials are retried.
+        # Test max_retry=1 to see if trial is retried.
+        assert RetryFailedTrialCallback.retried_trial_number(study.trials[1]) == (
+            None if max_retry == 0 else 0
+        )


Code review comments:
diff_hunk: @@ -607,3 +643,48 @@ def _check_trial_id(self, trial_id: int) -> None:
 
         if not self._redis.exists(self._key_trial(trial_id)):
             raise KeyError("study_id {} does not exist.".format(trial_id))
+
+    def record_heartbeat(self, trial_id: int) -> None:
+        study_id = self.get_study_id_from_trial_id(trial_id)
+        self._redis.hset(
+            self._key_study_heartbeats(study_id), str(trial_id), pickle.dumps(time.time())
path: optuna/storages/_redis.py
text: @user1:
Can you use the current time of DB side, not client side? c.f. https://github.com/optuna/optuna/blob/master/optuna/storages/_rdb/storage.py#L1197

diff_hunk: @@ -607,3 +643,48 @@ def _check_trial_id(self, trial_id: int) -> None:
 
         if not self._redis.exists(self._key_trial(trial_id)):
             raise KeyError("study_id {} does not exist.".format(trial_id))
+
+    def record_heartbeat(self, trial_id: int) -> None:
+        study_id = self.get_study_id_from_trial_id(trial_id)
+        self._redis.hset(
+            self._key_study_heartbeats(study_id), str(trial_id), pickle.dumps(time.time())
+        )
+
+    def fail_stale_trials(self, study_id: int) -> List[int]:
+        confirmed = []
+        for trial_id in self._get_stale_trial_ids(study_id):
+            if self.set_trial_state(trial_id, TrialState.FAIL):
+                confirmed.append(trial_id)
+
+        return confirmed
+
+    def _get_stale_trial_ids(self, study_id: int) -> List[int]:
+        assert self._heartbeat_interval is not None
+
+        if self._grace_period is None:
+            grace_period = 2 * self._heartbeat_interval
+        else:
+            grace_period = self._grace_period
+
+        current_time = time.time()
path: optuna/storages/_redis.py
text: @user1:
Can you use the current time of DB side, not client side?

@author:
Thanks, I was not aware of this problem. Good catch!