diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py b/rdagent/scenarios/data_science/proposal/exp_gen/merge.py index e541f159e8..1f12da20a6 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/merge.py @@ -1,6 +1,7 @@ """Merge the version in different traces""" import json +from contextlib import contextmanager from datetime import timedelta from typing import Dict, Tuple @@ -23,6 +24,19 @@ ) +@contextmanager +def _disable_auto_recovery(): + prev_t = DS_RD_SETTING.coding_fail_reanalyze_threshold + prev_e = DS_RD_SETTING.consecutive_errors + DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000 + DS_RD_SETTING.consecutive_errors = 100000 + try: + yield + finally: + DS_RD_SETTING.coding_fail_reanalyze_threshold = prev_t + DS_RD_SETTING.consecutive_errors = prev_e + + class MergeExpGen(ExpGen): def gen( self, @@ -262,14 +276,11 @@ 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) - else: - return self.merge_exp_gen.gen(trace) + with _disable_auto_recovery(): + if trace.sub_trace_count < 2: + return self.exp_gen.gen(trace) + else: + return self.merge_exp_gen.gen(trace) class MergeExpGen_MultiTrace(ExpGen): @@ -392,23 +403,20 @@ 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) - 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) - else: - # return self.merge_exp_gen.gen(trace) + with _disable_auto_recovery(): + leaves: list[int] = trace.get_leaves() + if len(leaves) < 2: trace.set_current_selection(selection=(-1,)) - return self.exp_gen.gen(trace) # continue the last trace, to polish the merged solution + return 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) + 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 class ExpGen2TraceAndMergeV3(ExpGen): @@ -428,20 +436,17 @@ 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) - else: - selection = (leaves[0],) - if trace.sota_exp_to_submit is not None: - for i in range(1, len(leaves)): - if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]): - selection = (leaves[i],) - break - trace.set_current_selection(selection) - return self.merge_exp_gen.gen(trace) + with _disable_auto_recovery(): + leaves: list[int] = trace.get_leaves() + if len(leaves) < 2: + trace.set_current_selection(selection=(-1,)) + return self.exp_gen.gen(trace) + else: + selection = (leaves[0],) + if trace.sota_exp_to_submit is not None: + for i in range(1, len(leaves)): + if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]): + selection = (leaves[i],) + break + trace.set_current_selection(selection) + return self.merge_exp_gen.gen(trace) 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..8cf96fe3f4 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING @@ -33,6 +34,19 @@ from rdagent.utils.workflow.loop import LoopBase +@contextmanager +def _disable_auto_recovery(): + prev_t = DS_RD_SETTING.coding_fail_reanalyze_threshold + prev_e = DS_RD_SETTING.consecutive_errors + DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000 + DS_RD_SETTING.consecutive_errors = 100000 + try: + yield + finally: + DS_RD_SETTING.coding_fail_reanalyze_threshold = prev_t + DS_RD_SETTING.consecutive_errors = prev_e + + class ParallelMultiTraceExpGen(ExpGen): """ An experiment generation strategy that enables parallel multi-trace exploration. @@ -115,10 +129,9 @@ 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_gen_type = type(self.merge_exp_gen).__name__ + with _disable_auto_recovery(): + exp = self.merge_exp_gen.gen(trace, plan=ds_plan) + 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 exp = self.exp_gen.gen(trace, plan=ds_plan)