From 8321cab98f5569866447348924949760e94a7625 Mon Sep 17 00:00:00 2001 From: "Bowen Xian (Shanghai Wicresoft Co Ltd)" Date: Thu, 16 Jul 2026 18:03:13 +0800 Subject: [PATCH] fix: scope merge auto-recovery state per experiment Replace process-wide threshold mutations with an experiment-level merge-phase marker so parallel traces cannot leak recovery state across iterations. Add regression coverage for all merge generators. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../data_science/experiment/experiment.py | 5 + rdagent/scenarios/data_science/loop.py | 3 +- .../data_science/proposal/exp_gen/merge.py | 32 +++---- .../proposal/exp_gen/router/__init__.py | 3 +- test/data_science/test_merge_auto_recovery.py | 96 +++++++++++++++++++ 5 files changed, 117 insertions(+), 22 deletions(-) create mode 100644 test/data_science/test_merge_auto_recovery.py diff --git a/rdagent/scenarios/data_science/experiment/experiment.py b/rdagent/scenarios/data_science/experiment/experiment.py index 209effdae7..a8bb0fb376 100644 --- a/rdagent/scenarios/data_science/experiment/experiment.py +++ b/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 diff --git a/rdagent/scenarios/data_science/loop.py b/rdagent/scenarios/data_science/loop.py index bdf028312a..1bd876b2dd 100644 --- a/rdagent/scenarios/data_science/loop.py +++ b/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: diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py b/rdagent/scenarios/data_science/proposal/exp_gen/merge.py index e541f159e8..f30bc86fc7 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py +++ b/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 diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py b/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py index 13e942af52..9e75b16d4d 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py +++ b/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 diff --git a/test/data_science/test_merge_auto_recovery.py b/test/data_science/test_merge_auto_recovery.py new file mode 100644 index 0000000000..e003777a7a --- /dev/null +++ b/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