fix: scope merge auto-recovery state per experiment

rdagent/scenarios/data_science/experiment/experiment.py

@@ -9,6 +9,8 @@
 
 
 class DSExperiment(Experiment[Task, FBWorkspace, FBWorkspace]):
+    is_merge_phase: bool = False
+
     def __init__(self, pending_tasks_list: list, hypothesis_candidates: list | None = None, *args, **kwargs) -> None:
         super().__init__(sub_tasks=[], *args, **kwargs)
         # Status
@@ -41,3 +43,6 @@ def is_ready_to_run(self) -> bool:
 
     def set_local_selection(self, local_selection: tuple[int, ...]) -> None:
         self.local_selection = local_selection
+
+    def set_merge_phase(self) -> None:
+        self.is_merge_phase = True

rdagent/scenarios/data_science/loop.py

@@ -211,6 +211,7 @@ def feedback(self, prev_out: dict[str, Any]) -> ExperimentFeedback:
     def record(self, prev_out: dict[str, Any]):
 
         exp: DSExperiment = None
+        generated_exp: DSExperiment = prev_out["direct_exp_gen"]
 
         cur_loop_id = prev_out[self.LOOP_IDX_KEY]
 
@@ -252,7 +253,7 @@ def record(self, prev_out: dict[str, Any]):
             )
             # Value backpropagation is handled in async_gen before next() via observe_commits
 
-            if self.trace.sota_experiment() is None:
+            if not generated_exp.is_merge_phase and self.trace.sota_experiment() is None:
                 if DS_RD_SETTING.coder_on_whole_pipeline:
                     #  check if feedback is not generated
                     if len(self.trace.hist) >= DS_RD_SETTING.coding_fail_reanalyze_threshold:

rdagent/scenarios/data_science/proposal/exp_gen/merge.py

@@ -262,14 +262,12 @@ def gen(
             trace.set_current_selection(selection)
             return self.exp_gen.gen(trace)
         else:
-            # disable reset in merging stage
-            DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
-            DS_RD_SETTING.consecutive_errors = 100000
-
             if trace.sub_trace_count < 2:
-                return self.exp_gen.gen(trace)
+                exp = self.exp_gen.gen(trace)
             else:
-                return self.merge_exp_gen.gen(trace)
+                exp = self.merge_exp_gen.gen(trace)
+            exp.set_merge_phase()
+            return exp
 
 
 class MergeExpGen_MultiTrace(ExpGen):
@@ -392,23 +390,21 @@ def gen(
             return self.exp_gen.gen(trace)
 
         else:
-            # disable reset in merging stage
-            DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
-            DS_RD_SETTING.consecutive_errors = 100000
-
             leaves: list[int] = trace.get_leaves()
             if len(leaves) < 2:
                 trace.set_current_selection(selection=(-1,))
-                return self.exp_gen.gen(trace)
+                exp = self.exp_gen.gen(trace)
             else:
                 if not self.flag_start_merge:  # root node of the merge trace
                     self.flag_start_merge = True
                     trace.set_current_selection(trace.NEW_ROOT)
-                    return self.merge_exp_gen.gen(trace)
+                    exp = self.merge_exp_gen.gen(trace)
                 else:
                     # return self.merge_exp_gen.gen(trace)
                     trace.set_current_selection(selection=(-1,))
-                    return self.exp_gen.gen(trace)  # continue the last trace, to polish the merged solution
+                    exp = self.exp_gen.gen(trace)  # continue the last trace, to polish the merged solution
+            exp.set_merge_phase()
+            return exp
 
 
 class ExpGen2TraceAndMergeV3(ExpGen):
@@ -428,14 +424,10 @@ def gen(
         if timer.remain_time() >= timedelta(hours=DS_RD_SETTING.merge_hours):
             return self.exp_gen.gen(trace)
         else:
-            # disable reset in merging stage
-            DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
-            DS_RD_SETTING.consecutive_errors = 100000
-
             leaves: list[int] = trace.get_leaves()
             if len(leaves) < 2:
                 trace.set_current_selection(selection=(-1,))
-                return self.exp_gen.gen(trace)
+                exp = self.exp_gen.gen(trace)
             else:
                 selection = (leaves[0],)
                 if trace.sota_exp_to_submit is not None:
@@ -444,4 +436,6 @@ def gen(
                             selection = (leaves[i],)
                             break
                 trace.set_current_selection(selection)
-                return self.merge_exp_gen.gen(trace)
+                exp = self.merge_exp_gen.gen(trace)
+            exp.set_merge_phase()
+            return exp

rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py

@@ -115,9 +115,8 @@ async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
                     and timer.remain_time() < timedelta(hours=DS_RD_SETTING.merge_hours)
                     and len(leaves) >= 2
                 ):
-                    DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
-                    DS_RD_SETTING.consecutive_errors = 100000
                     exp = self.merge_exp_gen.gen(trace, plan=ds_plan)
+                    exp.set_merge_phase()
                     exp_gen_type = type(self.merge_exp_gen).__name__
                 else:
                     # If there is a sota experiment in the sub-trace and not in merge time, we use default exp_gen

test/data_science/test_merge_auto_recovery.py

@@ -0,0 +1,96 @@
+import asyncio
+from datetime import timedelta
+from unittest.mock import Mock, patch
+
+import pytest
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
+from rdagent.scenarios.data_science.proposal.exp_gen.merge import (
+    ExpGen2TraceAndMerge,
+    ExpGen2TraceAndMergeV2,
+    ExpGen2TraceAndMergeV3,
+)
+from rdagent.scenarios.data_science.proposal.exp_gen.router import ParallelMultiTraceExpGen
+
+
+def make_experiment() -> DSExperiment:
+    return DSExperiment.__new__(DSExperiment)
+
+
+@pytest.mark.offline
+@pytest.mark.parametrize("generator_class", [ExpGen2TraceAndMerge, ExpGen2TraceAndMergeV2, ExpGen2TraceAndMergeV3])
+def test_merge_generators_use_per_experiment_state(generator_class: type) -> None:
+    experiment = make_experiment()
+    merge_exp_gen = Mock()
+    merge_exp_gen.gen.return_value = experiment
+    generator = object.__new__(generator_class)
+    generator.exp_gen = Mock()
+    generator.merge_exp_gen = merge_exp_gen
+    if generator_class is ExpGen2TraceAndMergeV2:
+        generator.flag_start_merge = False
+
+    trace = Mock()
+    trace.NEW_ROOT = ()
+    trace.sota_exp_to_submit = None
+    trace.get_leaves.return_value = [0, 1]
+    trace.sub_trace_count = 2
+
+    original_thresholds = (
+        DS_RD_SETTING.coding_fail_reanalyze_threshold,
+        DS_RD_SETTING.consecutive_errors,
+    )
+    with patch(
+        "rdagent.scenarios.data_science.proposal.exp_gen.merge.RD_Agent_TIMER_wrapper.timer",
+    ) as timer:
+        timer.remain_time.return_value = timedelta()
+        result = generator.gen(trace)
+
+    assert result.is_merge_phase is True
+    assert (
+        DS_RD_SETTING.coding_fail_reanalyze_threshold,
+        DS_RD_SETTING.consecutive_errors,
+    ) == original_thresholds
+
+
+@pytest.mark.offline
+def test_parallel_merge_generator_uses_per_experiment_state() -> None:
+    experiment = make_experiment()
+    generator = object.__new__(ParallelMultiTraceExpGen)
+    generator.merge_exp_gen = Mock()
+    generator.merge_exp_gen.gen.return_value = experiment
+    generator.exp_gen = Mock()
+    generator.draft_exp_gen = Mock()
+    generator.planner = Mock()
+    generator.trace_scheduler = Mock()
+
+    trace = Mock()
+    trace.get_leaves.return_value = [0, 1]
+    trace.sota_exp_to_submit = None
+    loop = Mock()
+    loop.loop_idx = 1
+    loop.get_unfinished_loop_cnt.return_value = 0
+
+    original_thresholds = (
+        DS_RD_SETTING.coding_fail_reanalyze_threshold,
+        DS_RD_SETTING.consecutive_errors,
+    )
+    with (
+        patch(
+            "rdagent.scenarios.data_science.proposal.exp_gen.router.RD_Agent_TIMER_wrapper.timer",
+        ) as timer,
+        patch.object(DS_RD_SETTING, "enable_planner", new=False),
+    ):
+        timer.started = True
+        timer.remain_time.return_value = timedelta()
+        result = asyncio.run(generator.async_gen(trace, loop))
+
+    assert result.is_merge_phase is True
+    assert (
+        DS_RD_SETTING.coding_fail_reanalyze_threshold,
+        DS_RD_SETTING.consecutive_errors,
+    ) == original_thresholds
+
+
+@pytest.mark.offline
+def test_regular_experiment_enables_auto_recovery_by_default() -> None:
+    assert make_experiment().is_merge_phase is False