From 8bdccb59d70242768231ae8d270ea56648f6867f Mon Sep 17 00:00:00 2001 From: you-n-g Date: Tue, 18 Feb 2025 21:56:32 +0800 Subject: [PATCH] fix: move next_component_required logic to DSTrace class and accurate implement (#612) * refactor: Move next_component_required logic to DSTrace class * lint --- rdagent/app/data_science/loop.py | 18 ++++++--------- .../data_science/experiment/experiment.py | 14 ------------ .../data_science/proposal/exp_gen.py | 22 ++++++++++++++----- 3 files changed, 24 insertions(+), 30 deletions(-) diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index f248dbc5..93fae288 100644 --- a/rdagent/app/data_science/loop.py +++ b/rdagent/app/data_science/loop.py @@ -93,7 +93,7 @@ class DataScienceRDLoop(RDLoop): def running(self, prev_out: dict[str, Any]): exp: DSExperiment = prev_out["coding"] - if exp.next_component_required() is None: + if self.trace.next_component_required() is None: new_exp = self.runner.develop(exp) logger.log_object(new_exp) return new_exp @@ -102,7 +102,7 @@ class DataScienceRDLoop(RDLoop): def feedback(self, prev_out: dict[str, Any]) -> ExperimentFeedback: exp: DSExperiment = prev_out["running"] - if exp.next_component_required() is None: + if self.trace.next_component_required() is None: feedback = self.summarizer.generate_feedback(exp, self.trace) else: feedback = ExperimentFeedback( @@ -124,15 +124,11 @@ class DataScienceRDLoop(RDLoop): ) ) if self.trace.sota_experiment() is None and len(self.trace.hist) >= DS_RD_SETTING.consecutive_errors: - trace_exp_next_component_list = [ - type(exp.pending_tasks_list[0][0]) - for exp, _ in self.trace.hist[-DS_RD_SETTING.consecutive_errors :] - ] - last_successful_exp = self.trace.last_successful_exp() - if ( - last_successful_exp not in [exp for exp, _ in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]] - and len(set(trace_exp_next_component_list)) == 1 - ): + # if {in inital/drafting stage} and {tried enough times} + for _, fb in self.trace.hist[-DS_RD_SETTING.consecutive_errors :]: + if fb: + break # any success will stop restarting. + else: # otherwise restart it logger.error("Consecutive errors reached the limit. Dumping trace.") logger.log_object(self.trace, tag="trace before restart") self.trace = DSTrace(scen=self.trace.scen, knowledge_base=self.trace.knowledge_base) diff --git a/rdagent/scenarios/data_science/experiment/experiment.py b/rdagent/scenarios/data_science/experiment/experiment.py index 68fe18e6..d071807d 100644 --- a/rdagent/scenarios/data_science/experiment/experiment.py +++ b/rdagent/scenarios/data_science/experiment/experiment.py @@ -19,17 +19,3 @@ class DSExperiment(Experiment[Task, FBWorkspace, FBWorkspace]): self.experiment_workspace = FBWorkspace() self.pending_tasks_list = pending_tasks_list self.format_check_result = None - - def next_component_required(self) -> COMPONENT | None: - files = list(self.experiment_workspace.file_dict.keys()) - if "load_data.py" not in files: - return "DataLoadSpec" - if "feature.py" not in files: - return "FeatureEng" - if not any(re.match(r"model.*\.py", file) for file in files): - return "Model" - if "ensemble.py" not in files: - return "Ensemble" - if "main.py" not in files: - return "Workflow" - return None diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py index 2bf225d7..16fef349 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen.py @@ -90,6 +90,21 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): self.hist: list[tuple[DSExperiment, ExperimentFeedback]] = [] self.knowledge_base = knowledge_base + COMPLETE_ORDER = ("DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow") + + def next_component_required(self) -> COMPONENT | None: + for c in self.COMPLETE_ORDER: + if not self.has_compponent(c): + return c + return None + + def has_compponent(self, component: COMPONENT) -> bool: + for exp, fb in self.hist: + assert isinstance(exp.hypothesis, DSHypothesis), "Hypothesis should be DSHypothesis (and not None)" + if exp.hypothesis.component == component and fb: + return True + return False + def sota_experiment(self, last_n: int = -1) -> DSExperiment | None: """ Access the last experiment result. @@ -108,7 +123,7 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): assert last_n < 0 for exp, ef in self.hist[::-1]: # the sota exp should be accepted decision and all required components are completed. - if ef.decision and exp.next_component_required() is None: + if ef.decision and self.next_component_required() is None: last_n += 1 if last_n == 0: return exp @@ -237,10 +252,7 @@ class DSExpGen(ExpGen): scenario_desc = trace.scen.get_scenario_all_desc() last_successful_exp = trace.last_successful_exp() - if len(trace.hist) == 0 or last_successful_exp is None: - next_missing_component = "DataLoadSpec" - else: - next_missing_component = last_successful_exp.next_component_required() + next_missing_component = trace.next_component_required() init_component_config = { "DataLoadSpec": {"task_cls": DataLoaderTask, "spec_file": None, "component_prompt_key": "data_loader"},