Merge pull request #32 from xadrianzetx/task-dedupe-v2

Task state based deduplication
xadrianzetx · Oct 19, 2022 · 8101460 · 8101460
2 parents c4aa7f8 + d37f636
commit 8101460
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 46 deletions.
diff --git a/optuna_distributed/managers/distributed.py b/optuna_distributed/managers/distributed.py
@@ -200,8 +200,10 @@ def register_trial_exit(self, trial_id: int) -> None:
 
 def _distributable(func: ObjectiveFuncType, with_supervisor: bool) -> DistributableWithContext:
     def _wrapper(context: _TaskContext) -> None:
-        # FIXME: Re-introduce task deduplication.
         task_state = Variable(context.state_id)
+        if task_state.get() != _TaskState.WAITING:
+            return
+
         task_state.set(_TaskState.RUNNING)
         message: Message
 

diff --git a/optuna_distributed/messages/__init__.py b/optuna_distributed/messages/__init__.py
@@ -5,7 +5,6 @@
 from optuna_distributed.messages.property import TrialProperty
 from optuna_distributed.messages.property import TrialPropertyMessage
 from optuna_distributed.messages.pruned import PrunedMessage
-from optuna_distributed.messages.repeated import RepeatedTrialMessage
 from optuna_distributed.messages.report import ReportMessage
 from optuna_distributed.messages.response import ResponseMessage
 from optuna_distributed.messages.setattr import AttributeType
@@ -21,7 +20,6 @@
     "SuggestMessage",
     "CompletedMessage",
     "FailedMessage",
-    "RepeatedTrialMessage",
     "PrunedMessage",
     "ReportMessage",
     "ShouldPruneMessage",

diff --git a/optuna_distributed/messages/repeated.py b/optuna_distributed/messages/repeated.py
diff --git a/tests/test_managers.py b/tests/test_managers.py
@@ -2,6 +2,8 @@
 import multiprocessing
 import sys
 import time
+from unittest.mock import Mock
+import uuid
 
 from dask.distributed import Client
 from dask.distributed import Variable
@@ -13,7 +15,9 @@
 from optuna_distributed.managers import LocalOptimizationManager
 from optuna_distributed.managers import ObjectiveFuncType
 from optuna_distributed.managers.distributed import _StateSynchronizer
+from optuna_distributed.managers.distributed import _TaskContext
 from optuna_distributed.managers.distributed import _TaskState
+from optuna_distributed.managers.distributed import _distributable
 from optuna_distributed.messages import CompletedMessage
 from optuna_distributed.messages import HeartbeatMessage
 from optuna_distributed.messages import ResponseMessage
@@ -117,6 +121,27 @@ def _objective(trial: DistributedTrial) -> float:
             break
 
 
+def test_distributed_task_deduped(client: Client) -> None:
+    def _objective(trial: DistributedTrial) -> float:
+        run_count = Variable("run_count")
+        run_count.set(run_count.get() + 1)
+        return 0.0
+
+    run_count = Variable("run_count")
+    run_count.set(0)
+    state_id = uuid.uuid4().hex
+    Variable(state_id).set(_TaskState.WAITING)
+
+    # Simulate scenario where task run was repeated.
+    # https://stackoverflow.com/a/41965766
+    func = _distributable(_objective, with_supervisor=False)
+    context = _TaskContext(DistributedTrial(0, Mock()), stop_flag="foo", state_id=state_id)
+    for _ in range(5):
+        client.submit(func, context).result()
+
+    assert run_count.get() == 1
+
+
 def test_synchronizer_optimization_enabled() -> None:
     synchronizer = _StateSynchronizer()
     optimization_enabled = Variable(synchronizer.stop_flag)

diff --git a/tests/test_messages.py b/tests/test_messages.py
@@ -21,7 +21,6 @@
 from optuna_distributed.messages import HeartbeatMessage
 from optuna_distributed.messages import Message
 from optuna_distributed.messages import PrunedMessage
-from optuna_distributed.messages import RepeatedTrialMessage
 from optuna_distributed.messages import ReportMessage
 from optuna_distributed.messages import ResponseMessage
 from optuna_distributed.messages import SetAttributeMessage
@@ -179,15 +178,6 @@ def test_should_prune(study: Study, manager: MockOptimizationManager) -> None:
     assert trial[0]._trial_id == 0
 
 
-def test_repeated_trial(study: Study, manager: MockOptimizationManager) -> None:
-    msg = RepeatedTrialMessage(0)
-    assert not msg.closing
-
-    study.tell(0, state=TrialState.PRUNED)
-    msg.process(study, manager)
-    assert _message_responds_with(True, manager=manager)
-
-
 def test_report_intermediate(study: Study, manager: MockOptimizationManager) -> None:
     msg = ReportMessage(0, value=0.0, step=1)
     assert not msg.closing