Fault tolerance: heartbeats and retries

On spot instances and shared clusters workers die mid-trial; without heartbeats those trials hang as RUNNING forever.

When a worker process is killed (OOM, preemption, a node failure), its trial never reports back. Without extra configuration it stays RUNNING indefinitely: it pollutes counts, MaxTrialsCallback miscounts, and the sampler may treat it as still in progress.

Heartbeats solve this for RDB storage:

storage = optuna.storages.RDBStorage(
    url, heartbeat_interval=60, grace_period=120,
    heartbeat_stale_trial_callback=RetryHeartbeatStaleTrialCallback(max_retry=3))

Each running trial records a heartbeat every heartbeat_interval seconds from a background thread.
A trial whose last heartbeat is older than grace_period is stale. Other workers check for stale trials when they start a new trial (or you call optuna.storages.fail_stale_trials(study)), and mark them FAIL.
RetryHeartbeatStaleTrialCallback(max_retry=3) then re-enqueues the same parameters as a WAITING trial so the configuration is retried. RetryHeartbeatStaleTrialCallback.retried_trial_number(trial) and retry_history(trial) link retries to the original. (It replaces the older RetryFailedTrialCallback, deprecated in 4.9.)

Other robustness practices:

Make objectives idempotent and checkpointed: save model checkpoints keyed by trial number so a retry can resume, and store paths in user attributes.
Treat expected failures inside the objective (catch=, returning NaN) separately from infrastructure failures (heartbeats).
Keep the database healthy: connection pooling via engine_kwargs, a gRPC proxy for hundreds of workers, regular backups.
Clean up abandoned studies with optuna.delete_study.

Example

import multiprocessing as mp
import time
import optuna
from optuna.storages import RDBStorage, RetryHeartbeatStaleTrialCallback

def make_storage():
    return RDBStorage("sqlite:///ft_demo.db", heartbeat_interval=1, grace_period=3,
        heartbeat_stale_trial_callback=RetryHeartbeatStaleTrialCallback(max_retry=2))

def worker():                                 # a trial that never finishes
    study = optuna.load_study(study_name="ft", storage=make_storage())
    study.optimize(lambda t: time.sleep(t.suggest_float("x", 0, 1) + 60), n_trials=1)

if __name__ == "__main__":
    optuna.logging.set_verbosity(optuna.logging.WARNING)
    study = optuna.create_study(study_name="ft", storage=make_storage())
    p = mp.Process(target=worker); p.start(); time.sleep(4)
    p.kill()                                  # simulate preemption mid-trial
    time.sleep(4)                             # heartbeat goes stale
    optuna.storages.fail_stale_trials(study)  # other workers do this automatically
    print([t.state.name for t in study.get_trials()])   # ['FAIL', 'WAITING']
    retry = study.get_trials()[-1]
    print(RetryHeartbeatStaleTrialCallback.retried_trial_number(retry))   # 0
    optuna.delete_study(study_name="ft", storage=make_storage())

Remember

Without heartbeats, crashed trials stay RUNNING forever
Stale trials (no heartbeat within grace_period) are marked FAIL by other workers
RetryHeartbeatStaleTrialCallback re-queues their parameters as WAITING trials

Watch out for

Setting grace_period shorter than the heartbeat interval plus normal delays, so healthy trials get marked stale.

Check yourself

A preempted worker leaves a trial RUNNING. With heartbeats and RetryHeartbeatStaleTrialCallback configured, what eventually happens?

Official docs: optuna.readthedocs.io/en/stable/reference/generated/optuna.storages.RDBStorage.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks