Merge pull request #5431 from not522/fix-heartbeat

Fix heartbeat for race condition
optuna · May 8, 2024 · ab958d4 · ab958d4
2 parents 8762a2f + 905ccd7
commit ab958d4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optuna/storages/_rdb/storage.py b/optuna/storages/_rdb/storage.py
@@ -4,6 +4,7 @@
 from contextlib import contextmanager
 import copy
 from datetime import datetime
+from datetime import timedelta
 import json
 import logging
 import os
@@ -1010,7 +1011,7 @@ def _get_stale_trial_ids(self, study_id: int) -> List[int]:
                     continue
                 assert len(trial.heartbeats) == 1
                 heartbeat = trial.heartbeats[0].heartbeat
-                if (current_heartbeat - heartbeat).seconds > grace_period:
+                if current_heartbeat - heartbeat > timedelta(seconds=grace_period):
                     stale_trial_ids.append(trial.trial_id)
 
         return stale_trial_ids