Repository: optuna/optuna

Issue:
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.



Discussion:

Resolving patch:
diff --git a/optuna/storages/_redis.py b/optuna/storages/_redis.py
--- a/optuna/storages/_redis.py
+++ b/optuna/storages/_redis.py
@@ -2,6 +2,7 @@
 from datetime import datetime
 import pickle
 from typing import Any
+from typing import Callable
 from typing import cast
 from typing import Dict
 from typing import List
@@ -59,19 +60,51 @@ 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.trial.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,11 @@ 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)
+
+            # To ensure that there are no failed trials with heartbeats in the DB
+            # under any circumstances
+            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 +645,54 @@ 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(self._get_redis_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 = self._get_redis_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 _get_redis_time(self) -> float:
+        seconds, microseconds = self._redis.time()
+        return seconds + microseconds * 1e-6
+
+    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


Test patch:
diff --git a/optuna/testing/storage.py b/optuna/testing/storage.py
--- 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
--- 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
--- 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
+        )