From fd155d1fa85a205fddb21a3239fda336fb5f44c9 Mon Sep 17 00:00:00 2001 From: xuangu-fang Date: Wed, 9 Apr 2025 09:42:30 +0800 Subject: [PATCH] feat: checkpoint selection (#744) * rebase selection code * bug-free run: checkpoint selection and dynamic EDA loading * add prototypes of various selectors, to imp. and test later * fix EDA write bug * move selector to from proposal.py tp seletc.py * auto lint * fix line-too-long typos * aligh the design of "selection", rm extra instance check * make auto-lint * add non-trival selector: SOTAjump --- rdagent/app/data_science/loop.py | 12 +- .../coder/data_science/ensemble/__init__.py | 2 +- .../coder/data_science/feature/__init__.py | 2 +- .../coder/data_science/model/__init__.py | 2 +- .../coder/data_science/pipeline/__init__.py | 2 +- .../coder/data_science/pipeline/eval.py | 4 +- .../data_science/raw_data_loader/__init__.py | 8 +- .../coder/data_science/workflow/__init__.py | 2 +- .../coder/data_science/workflow/eval.py | 3 +- rdagent/core/proposal.py | 32 +++- .../scenarios/data_science/dev/feedback.py | 5 +- .../scenarios/data_science/dev/runner/eval.py | 2 +- .../data_science/proposal/exp_gen/__init__.py | 6 +- .../data_science/proposal/exp_gen/base.py | 169 +++++++++++++++--- .../data_science/proposal/exp_gen/draft.py | 13 +- .../data_science/proposal/exp_gen/proposal.py | 19 +- .../data_science/proposal/exp_gen/select.py | 111 ++++++++++++ .../scenarios/data_science/scen/__init__.py | 8 +- 18 files changed, 353 insertions(+), 49 deletions(-) create mode 100644 rdagent/scenarios/data_science/proposal/exp_gen/select.py diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index e7c9f6ef..c12af725 100644 --- a/rdagent/app/data_science/loop.py +++ b/rdagent/app/data_science/loop.py @@ -27,6 +27,10 @@ from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner from rdagent.scenarios.data_science.experiment.experiment import DSExperiment from rdagent.scenarios.data_science.proposal.exp_gen import DSExpGen, DSTrace +from rdagent.scenarios.data_science.proposal.exp_gen.select import ( + LatestCKPSelector, + SOTAJumpCKPSelector, +) from rdagent.scenarios.kaggle.kaggle_crawler import download_data @@ -49,6 +53,7 @@ class DataScienceRDLoop(RDLoop): # 2) task generation from a complete solution # self.exp_gen: ExpGen = import_class(PROP_SETTING.exp_gen)(scen) + self.ckp_selector = LatestCKPSelector() self.exp_gen = DSExpGen(scen) self.data_loader_coder = DataLoaderCoSTEER(scen) self.feature_coder = FeatureCoSTEER(scen) @@ -68,7 +73,8 @@ class DataScienceRDLoop(RDLoop): super(RDLoop, self).__init__() def direct_exp_gen(self, prev_out: dict[str, Any]): - exp = self.exp_gen.gen(self.trace) + selection = self.ckp_selector.get_selection(self.trace) + exp = self.exp_gen.gen(self.trace, selection) logger.log_object(exp) # FIXME: this is for LLM debug webapp, remove this when the debugging is done. @@ -126,6 +132,10 @@ class DataScienceRDLoop(RDLoop): return feedback def record(self, prev_out: dict[str, Any]): + + # set the DAG parent for the trace + self.trace.sync_dag_parent_and_hist() + e = prev_out.get(self.EXCEPTION_KEY, None) if e is None: self.trace.hist.append((prev_out["running"], prev_out["feedback"])) diff --git a/rdagent/components/coder/data_science/ensemble/__init__.py b/rdagent/components/coder/data_science/ensemble/__init__.py index f25c5f56..e1882fa5 100644 --- a/rdagent/components/coder/data_science/ensemble/__init__.py +++ b/rdagent/components/coder/data_science/ensemble/__init__.py @@ -74,7 +74,7 @@ class EnsembleMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): ) # Generate code with knowledge integration - competition_info = self.scen.get_scenario_all_desc() + competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)) system_prompt = T(".prompts:ensemble_coder.system").r( task_desc=ensemble_information_str, competition_info=competition_info, diff --git a/rdagent/components/coder/data_science/feature/__init__.py b/rdagent/components/coder/data_science/feature/__init__.py index 38691aff..42b4359f 100644 --- a/rdagent/components/coder/data_science/feature/__init__.py +++ b/rdagent/components/coder/data_science/feature/__init__.py @@ -61,7 +61,7 @@ class FeatureMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): # 2. code system_prompt = T(".prompts:feature_coder.system").r( - competition_info=self.scen.get_scenario_all_desc(), + competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)), task_desc=feature_information_str, data_loader_code=workspace.file_dict.get("load_data.py"), queried_similar_successful_knowledge=queried_similar_successful_knowledge, diff --git a/rdagent/components/coder/data_science/model/__init__.py b/rdagent/components/coder/data_science/model/__init__.py index e35729e2..fa628f5b 100644 --- a/rdagent/components/coder/data_science/model/__init__.py +++ b/rdagent/components/coder/data_science/model/__init__.py @@ -62,7 +62,7 @@ class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): # 2. code system_prompt = T(".prompts:model_coder.system").r( task_desc=model_information_str, - competition_info=self.scen.get_scenario_all_desc(), + competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)), data_loader_code=workspace.file_dict.get("load_data.py"), feature_code=workspace.file_dict["feature.py"], queried_similar_successful_knowledge=queried_similar_successful_knowledge, diff --git a/rdagent/components/coder/data_science/pipeline/__init__.py b/rdagent/components/coder/data_science/pipeline/__init__.py index 73b003f8..b04cc0f5 100644 --- a/rdagent/components/coder/data_science/pipeline/__init__.py +++ b/rdagent/components/coder/data_science/pipeline/__init__.py @@ -66,7 +66,7 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): workspace: FBWorkspace | None = None, prev_task_feedback: CoSTEERSingleFeedback | None = None, ) -> dict[str, str]: - competition_info = self.scen.get_scenario_all_desc() + competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)) runtime_environment = self.scen.get_runtime_environment() data_folder_info = self.scen.processed_data_folder_description pipeline_task_info = target_task.get_task_information() diff --git a/rdagent/components/coder/data_science/pipeline/eval.py b/rdagent/components/coder/data_science/pipeline/eval.py index 0f4d6cff..809409fb 100644 --- a/rdagent/components/coder/data_science/pipeline/eval.py +++ b/rdagent/components/coder/data_science/pipeline/eval.py @@ -119,8 +119,10 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): ) stdout += "\n" + submission_check_out + eda_output = implementation.file_dict.get("EDA.md", None) + system_prompt = T(".prompts:pipeline_eval.system").r( - scenario=self.scen.get_scenario_all_desc(), + scenario=self.scen.get_scenario_all_desc(eda_output=eda_output), task_desc=target_task.get_task_information(), spec=T("scenarios.data_science.share:component_spec.Pipeline").r(), ) diff --git a/rdagent/components/coder/data_science/raw_data_loader/__init__.py b/rdagent/components/coder/data_science/raw_data_loader/__init__.py index 5f4c19f4..4f2a57e8 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/__init__.py +++ b/rdagent/components/coder/data_science/raw_data_loader/__init__.py @@ -67,7 +67,7 @@ class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): ) -> dict[str, str]: # return a workspace with "load_data.py", "spec/load_data.md" inside # assign the implemented code to the new workspace. - competition_info = self.scen.get_scenario_all_desc() + competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)) runtime_environment = self.scen.get_runtime_environment() data_folder_info = self.scen.processed_data_folder_description data_loader_task_info = target_task.get_task_information() @@ -231,5 +231,9 @@ class DataLoaderCoSTEER(CoSTEER): stdout = new_exp.experiment_workspace.execute(env=env, entry=f"python test/data_loader_test.py") match = re.search(r"(.*?)=== Start of EDA part ===(.*)=== End of EDA part ===", stdout, re.DOTALL) eda_output = match.groups()[1] if match else None - self.scen.eda_output = eda_output + if eda_output is not None: + new_exp.experiment_workspace.inject_files(**{"EDA.md": eda_output}) + else: + eda_output = "No EDA output." + new_exp.experiment_workspace.inject_files(**{"EDA.md": eda_output}) return new_exp diff --git a/rdagent/components/coder/data_science/workflow/__init__.py b/rdagent/components/coder/data_science/workflow/__init__.py index e63c543b..ba468f2d 100644 --- a/rdagent/components/coder/data_science/workflow/__init__.py +++ b/rdagent/components/coder/data_science/workflow/__init__.py @@ -59,7 +59,7 @@ class WorkflowMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): # 2. code system_prompt = T(".prompts:workflow_coder.system").r( task_desc=workflow_information_str, - competition_info=self.scen.get_scenario_all_desc(), + competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)), queried_similar_successful_knowledge=queried_similar_successful_knowledge, queried_former_failed_knowledge=queried_former_failed_knowledge[0], out_spec=PythonAgentOut.get_spec(), diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py index 555c34bb..c66b2956 100644 --- a/rdagent/components/coder/data_science/workflow/eval.py +++ b/rdagent/components/coder/data_science/workflow/eval.py @@ -127,7 +127,8 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): stdout += "\n" + submission_check_out system_prompt = T(".prompts:workflow_eval.system").r( - scenario=self.scen.get_scenario_all_desc(), + # here we pass `None` to `eda_output` because we do not have nor need EDA output for workflow. + scenario=self.scen.get_scenario_all_desc(eda_output=None), task_desc=target_task.get_task_information(), spec=( implementation.file_dict["spec/workflow.md"] diff --git a/rdagent/core/proposal.py b/rdagent/core/proposal.py index 5aba6e66..9441affa 100644 --- a/rdagent/core/proposal.py +++ b/rdagent/core/proposal.py @@ -1,4 +1,4 @@ -""" """ +# TODO: remove `self.scen` if traces will be passed into the instance. from __future__ import annotations @@ -112,9 +112,16 @@ ASpecificKB = TypeVar("ASpecificKB", bound=KnowledgeBase) class Trace(Generic[ASpecificScen, ASpecificKB]): + NodeType = tuple[Experiment, ExperimentFeedback] # Define NodeType as a new type representing the tuple + def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None: self.scen: ASpecificScen = scen - self.hist: list[tuple[Experiment, ExperimentFeedback]] = [] + self.hist: list[Trace.NodeType] = ( + [] + ) # List of tuples containing experiments and their feedback, organized over time. + self.dag_parent: list[tuple[int, ...]] = [] # List of tuples representing parent indices in the DAG structure. + # (,) represents no parent; (1,) presents one parent; (1, 2) represents two parents. + # TODO: self.hist is 2-tuple now, remove hypothesis from it, change old code for this later. self.knowledge_base: ASpecificKB | None = knowledge_base @@ -128,13 +135,32 @@ class Trace(Generic[ASpecificScen, ASpecificKB]): return None, None +class CheckpointSelector: + """ + In the trace, we may start from any check point (we'll represent it as a variable `from_checkpoint_idx`) + """ + + @abstractmethod + def get_selection(self, trace: Trace) -> tuple[int, ...] | None: + """ + checkpoint_idx represents the place where we want to create a new node. + the return value should be the idx of target node (the parent of the new generating node). + - `(-1, )` represents starting from the latest trial in the trace - default value + - `(idx, )` represents starting from the `idx`-th trial in the trace. + - `None` represents starting from scratch (start a new trace) + + + - More advanced selection strategies in `select.py` + """ + + class ExpGen(ABC): def __init__(self, scen: Scenario) -> None: self.scen = scen @abstractmethod - def gen(self, trace: Trace) -> Experiment: + def gen(self, trace: Trace, selection: tuple[int, ...] = (-1,)) -> Experiment: """ Generate the experiment based on the trace. diff --git a/rdagent/scenarios/data_science/dev/feedback.py b/rdagent/scenarios/data_science/dev/feedback.py index 89d9ea05..5dacae8d 100644 --- a/rdagent/scenarios/data_science/dev/feedback.py +++ b/rdagent/scenarios/data_science/dev/feedback.py @@ -86,7 +86,10 @@ class DSExperiment2Feedback(Experiment2Feedback): decision=False, ) - system_prompt = T(".prompts:exp_feedback.system").r(scenario=self.scen.get_scenario_all_desc()) + eda_output = exp.experiment_workspace.file_dict.get("EDA.md", None) + system_prompt = T(".prompts:exp_feedback.system").r( + scenario=self.scen.get_scenario_all_desc(eda_output=eda_output) + ) user_prompt = T(".prompts:exp_feedback.user").r( sota_desc=sota_desc, cur_exp=exp, diff --git a/rdagent/scenarios/data_science/dev/runner/eval.py b/rdagent/scenarios/data_science/dev/runner/eval.py index 308ed384..992a417f 100644 --- a/rdagent/scenarios/data_science/dev/runner/eval.py +++ b/rdagent/scenarios/data_science/dev/runner/eval.py @@ -124,7 +124,7 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): stdout += f"\nMLEBench submission check:\n{submission_check_out}\nIf MLEBench submission check returns a 'Submission is valid' or similar message, despite some warning messages, you should still consider the submission as valid and give a positive final decision. " system_prompt = T(".prompts:DSCoSTEER_eval.system").r( - scenario=self.scen.get_scenario_all_desc(), + scenario=self.scen.get_scenario_all_desc(eda_output=implementation.file_dict.get("EDA.md", None)), task_desc=target_task.get_task_information(), ) user_prompt = T(".prompts:DSCoSTEER_eval.user").r( diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/__init__.py b/rdagent/scenarios/data_science/proposal/exp_gen/__init__.py index fb806228..a6311cfd 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/__init__.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/__init__.py @@ -20,7 +20,11 @@ class DSExpGen(ExpGen): def __init__(self, scen: DataScienceScen) -> None: super().__init__(scen) - def gen(self, trace: DSTrace) -> DSExperiment: + def gen(self, trace: DSTrace, selection: tuple[int, ...] = (-1,)) -> DSExperiment: + + # set the current selection for the trace + # handy design:dynamically change the "current selection" attribute of the trace, and we donot need to pass selection as an argument to other functions + trace.set_current_selection(selection) if DS_RD_SETTING.proposal_version not in ["v1", "v2"]: return import_class(DS_RD_SETTING.proposal_version)(scen=self.scen).gen(trace=trace) diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/base.py b/rdagent/scenarios/data_science/proposal/exp_gen/base.py index 687349d4..ba129db3 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/base.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/base.py @@ -39,51 +39,147 @@ class DSHypothesis(Hypothesis): class DSTrace(Trace[DataScienceScen, KnowledgeBase]): + def __init__(self, scen: DataScienceScen, knowledge_base: KnowledgeBase | None = None) -> None: self.scen: DataScienceScen = scen self.hist: list[tuple[DSExperiment, ExperimentFeedback]] = [] + """ + The dag_parent is a list of tuples, each tuple is the parent index of the current node. + The first element of the tuple is the parent index, the rest are the parent indexes of the parent (not implemented yet). + If the current node is the root node without parent, the tuple is empty. + """ + self.dag_parent: list[tuple[int, ...]] = [] # List of tuples representing parent indices in the DAG structure. + # () represents no parent; (1,) presents one parent; (1, 2) represents two parents. + self.knowledge_base = knowledge_base + self.current_selection: tuple[int, ...] = (-1,) + COMPLETE_ORDER = ("DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow") - def next_incomplete_component(self) -> COMPONENT | None: + def get_current_selection(self) -> tuple[int, ...]: + return self.current_selection + + def set_current_selection(self, selection: tuple[int, ...]) -> None: + self.current_selection = selection + + def sync_dag_parent_and_hist( + self, + ) -> None: + """ + Adding corresponding parent index to the dag_parent when the hist is going to be changed. + Should be called when the hist is changed. + """ + + if len(self.hist) == 0 or len(self.get_current_selection()) == 0: + # the node we are going to add is the first node of hist / root node of a new sub-trace + self.dag_parent.append(()) + + else: + current_node_idx = self.current_selection[0] + + if current_node_idx == -1: + # the current selection is the latest one + current_node_idx = len(self.hist) - 1 + + self.dag_parent.append((current_node_idx,)) + + def retrieve_search_list( + self, search_type: Literal["all", "ancestors"] = "ancestors" + ) -> list[tuple[DSExperiment, ExperimentFeedback]]: + """ + Retrieve the search list based on the selection and search_type. + + Parameters + ---------- + search_type : str + One of "all", "ancestors". + - "all": search the whole hist. + - "ancestors": search the trace from root to the selection. + + Returns + ------- + list[tuple[DSExperiment, ExperimentFeedback]] + The search list. + """ + + selection = self.get_current_selection() + if selection is None: + # selection is None, which means we switch to a new trace, which is not implemented yet + return [] + + return self.collect_all_ancestors(selection) if search_type == "ancestors" else self.hist + + def collect_all_ancestors( + self, + selection: tuple[int, ...] = (-1,), + ) -> list[tuple[DSExperiment, ExperimentFeedback]]: + """ + Collect all ancestors of the given selection. + The return list follows the order of [root->...->parent->current_node]. + """ + + if len(self.dag_parent) == 0: + return [] + + else: + all_ancestors = [] + + # start from the latest selection + current_node_idx = selection[0] + + # add the current node to the list + all_ancestors.insert(0, self.hist[current_node_idx]) + + parent_idx = self.dag_parent[current_node_idx] + + while len(parent_idx) > 0: + all_ancestors.insert(0, self.hist[parent_idx[0]]) + parent_idx = self.dag_parent[parent_idx[0]] + + return all_ancestors + + def next_incomplete_component( + self, + search_type: Literal["all", "ancestors"] = "ancestors", + ) -> COMPONENT | None: """ NOTE: - A component will be complete until get True decision feedback !!! + """ + search_list = self.retrieve_search_list(search_type) + for c in self.COMPLETE_ORDER: - if not self.has_component(c): + """Check if the component is in the ancestors of the selection.""" + if not self.has_component(c, search_list): return c + return None - def has_component(self, component: COMPONENT) -> bool: - for exp, fb in self.hist: + def has_component( + self, component: COMPONENT, search_list: list[tuple[DSExperiment, ExperimentFeedback]] = [] + ) -> bool: + for exp, fb in search_list: assert isinstance(exp.hypothesis, DSHypothesis), "Hypothesis should be DSHypothesis (and not None)" if exp.hypothesis.component == component and fb: return True return False def experiment_and_feedback_list_after_init( - self, return_type: Literal["sota", "failed", "all"] + self, + return_type: Literal["sota", "failed", "all"], + search_type: Literal["all", "ancestors"] = "all", ) -> list[tuple[DSExperiment, ExperimentFeedback]]: """ Retrieve a list of experiments and feedbacks based on the return_type. - - Parameters - ---------- - return_type : str - One of "sota", "failed", "all". - - Returns - ------- - list[tuple[DSExperiment, ExperimentFeedback]] - List of experiments and feedbacks. """ + search_list = self.retrieve_search_list(search_type) final_component = self.COMPLETE_ORDER[-1] has_final_component = True if DS_RD_SETTING.coder_on_whole_pipeline else False exp_and_feedback_list = [] - for exp, fb in self.hist: + for exp, fb in search_list: if has_final_component: if return_type == "all": exp_and_feedback_list.append((exp, fb)) @@ -95,34 +191,63 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): has_final_component = True return exp_and_feedback_list - def sota_experiment(self) -> DSExperiment | None: + def sota_experiment( + self, + search_type: Literal["all", "ancestors"] = "ancestors", + ) -> DSExperiment | None: """ + Returns ------- Experiment or None The experiment result if found, otherwise None. """ + search_list = self.retrieve_search_list(search_type) + if DS_RD_SETTING.coder_on_whole_pipeline or self.next_incomplete_component() is None: - for exp, ef in self.hist[::-1]: + for exp, ef in search_list[::-1]: # the sota exp should be accepted decision and all required components are completed. if ef.decision: return exp return None - def last_successful_exp(self) -> DSExperiment | None: + def last_successful_exp( + self, + search_type: Literal["all", "ancestors"] = "ancestors", + ) -> DSExperiment | None: """ Access the last successful experiment even part of the components are not completed. """ - for exp, ef in self.hist[::-1]: + search_list = self.retrieve_search_list(search_type) + + for exp, ef in search_list[::-1]: if ef.decision: return exp return None - def last_runnable_exp_fb(self) -> tuple[DSExperiment, ExperimentFeedback] | None: + def last_exp( + self, + search_type: Literal["all", "ancestors"] = "ancestors", + ) -> DSExperiment | None: + """ + Access the last experiment + """ + search_list = self.retrieve_search_list(search_type) + + for exp, ef in search_list[::-1]: + return exp + return None + + def last_runnable_exp_fb( + self, + search_type: Literal["all", "ancestors"] = "ancestors", + ) -> tuple[DSExperiment, ExperimentFeedback] | None: """ Access the last runnable experiment (no exception, usually not all task failed) and feedback """ - for exp, ef in self.hist[::-1]: + search_list = self.retrieve_search_list(search_type) + + for exp, ef in search_list[::-1]: if ef.exception is None: return exp, ef return None diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/draft.py b/rdagent/scenarios/data_science/proposal/exp_gen/draft.py index bd1f6d4f..719ff8f0 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/draft.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/draft.py @@ -63,9 +63,15 @@ class DSDraftExpGen(ExpGen): scenario_desc: Description of the current scenario last_successful_exp: Last successful experiment or None spec_file: Path to specification file if needed + selection: The selection of the node to generate the task """ - scenario_desc = trace.scen.get_scenario_all_desc() last_successful_exp = trace.last_successful_exp() + # typecheck on the last successful exp, should be DSExperiment + if not isinstance(last_successful_exp, DSExperiment): + eda_output = None + else: + eda_output = last_successful_exp.experiment_workspace.file_dict.get("EDA.md", None) + scenario_desc = trace.scen.get_scenario_all_desc(eda_output=eda_output) init_component_config = { "DataLoadSpec": {"task_cls": DataLoaderTask, "spec_file": None, "component_prompt_key": "data_loader"}, "FeatureEng": {"task_cls": FeatureTask, "spec_file": "spec/feature.md", "component_prompt_key": "feature"}, @@ -78,8 +84,9 @@ class DSDraftExpGen(ExpGen): component_prompt_key = init_component_config[component].get("component_prompt_key") former_tasks_desc = "" - if len(trace.hist) > 0: - for exp, fb in reversed(trace.hist): + search_list = trace.retrieve_search_list() + if len(search_list) > 0: + for exp, fb in reversed(search_list): if exp is not last_successful_exp: former_task_desc = exp.pending_tasks_list[0][0].get_task_information() former_task_desc += f"\n\nYou have tried to implement the same component and got the following exception: \n{fb.exception}\n Please try different methods to avoid the same errors and results in an infinite loop" diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py index 75ed51aa..f522759d 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/proposal.py @@ -67,12 +67,17 @@ class DSProposalV1ExpGen(ExpGen): # - Previous Feedback # - Current sota implementation (encourage change based on it) # - Extra RAG - - scenario_desc = trace.scen.get_scenario_all_desc() sota_exp = trace.sota_experiment() + if not isinstance(sota_exp, DSExperiment): + eda_output = None + else: + eda_output = sota_exp.experiment_workspace.file_dict.get("EDA.md", None) + scenario_desc = trace.scen.get_scenario_all_desc(eda_output=eda_output) + assert sota_exp is not None, "SOTA experiment is not provided." - exp_and_feedback = trace.hist[-1] - last_exp = exp_and_feedback[0] + last_exp = trace.last_exp() + # exp_and_feedback = trace.hist[-1] + # last_exp = exp_and_feedback[0] # Step 1: Generate component # Describe current best solution using shared template @@ -395,7 +400,11 @@ class DSProposalV2ExpGen(ExpGen): ) sota_exp = trace.sota_experiment() - scenario_desc = trace.scen.get_scenario_all_desc() + if not isinstance(sota_exp, DSExperiment): + eda_output = None + else: + eda_output = sota_exp.experiment_workspace.file_dict.get("EDA.md", None) + scenario_desc = trace.scen.get_scenario_all_desc(eda_output=eda_output) competition_desc = trace.scen.get_competition_full_desc() sota_exp_desc = T("scenarios.data_science.share:describe.exp").r( diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/select.py b/rdagent/scenarios/data_science/proposal/exp_gen/select.py new file mode 100644 index 00000000..43b16e85 --- /dev/null +++ b/rdagent/scenarios/data_science/proposal/exp_gen/select.py @@ -0,0 +1,111 @@ +from rdagent.core.proposal import CheckpointSelector, Trace + +# # TODO: more advanced selector +# # TODO/Discussion: load selector function here or define selector class in `proposal.py`? + + +class LatestCKPSelector(CheckpointSelector): + """ + -`(-1, )` represents starting from the latest trial in the trace + """ + + def get_selection(self, trace: Trace) -> tuple[int, ...]: + + return (-1,) + + +class SOTAJumpCKPSelector(CheckpointSelector): + """ + SOTA jump policy: + if the cumulative SOTA in a window is below a threshold, jump to a new trial + otherwise, continue the current latest trial + """ + + def __init__(self) -> None: + + self.SOTA_COUNT_WINDOW = 5 + self.SOTA_COUNT_THRESHOLD = 1 # start to compute cumulative SOTA ratio after 5 trials + + def get_selection(self, trace: Trace) -> tuple[int, ...] | None: + + current_trace = trace.retrieve_search_list(search_type="ancestors") + if len(trace.hist) > 0 and len(current_trace) > self.SOTA_COUNT_WINDOW: + all_exp_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="ancestors") + # sota_exp_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="ancestors") + exp_list_in_window = all_exp_list[-self.SOTA_COUNT_WINDOW :] + + # compute the cumulative SOTA ratio in the window + sota_count = 0 + for exp, fb in exp_list_in_window: + if fb.decision: + sota_count += 1 + if sota_count < self.SOTA_COUNT_THRESHOLD: + return None + else: + return (-1,) + + else: + return (-1,) + + +# TODO: implement these selectors and more + + +class GlobalGreedyCKPSelector(CheckpointSelector): + """ + global greedy selector: select the trial with best performance globally (in trace.hist) + consistent with the greedy strategy in AIDE + not implemented yet + """ + + def get_selection(self, trace: Trace) -> tuple[int, ...]: + + return (-1,) + + +class LocalGreedyCKPSelector(CheckpointSelector): + """ + local greedy selector: select the trial with best performance locally (in trace.ancestors) + not implemented yet + """ + + def get_selection(self, trace: Trace) -> tuple[int, ...]: + + return (-1,) + + +class BugBufferCKPSelector(CheckpointSelector): + """ + bug buffer selector: with limit-size bug buffer size, start a new trace if buffer exceeds. + not implemented yet + """ + + def __init__(self) -> None: + self.bug_count = 0 + self.BUG_BUFFER_SIZE = 10 + + def get_selection(self, trace: Trace) -> tuple[int, ...]: + + if self.bug_count < self.BUG_BUFFER_SIZE: + return (-1,) + + else: + return None + + +class RandomCKPSelector(CheckpointSelector): + def get_selection(self, trace: Trace) -> tuple[int, ...]: + """ + random selector: select the trial randomly + not implemented yet + """ + return (-1,) + + +class BuggyCKPSelector(CheckpointSelector): + def get_selection(self, trace: Trace) -> tuple[int, ...]: + """ + buggy selector: select the most recent trial with buggy performance + not implemented yet + """ + return (-1,) diff --git a/rdagent/scenarios/data_science/scen/__init__.py b/rdagent/scenarios/data_science/scen/__init__.py index 05909839..15a3c609 100644 --- a/rdagent/scenarios/data_science/scen/__init__.py +++ b/rdagent/scenarios/data_science/scen/__init__.py @@ -32,7 +32,6 @@ class DataScienceScen(Scenario): self.processed_data_folder_description = self._get_data_folder_description() self._analysis_competition_description() self.metric_direction = self._get_direction() - self.eda_output = None def _get_description(self): if (fp := Path(f"{DS_RD_SETTING.local_data_path}/{self.competition}.json")).exists(): @@ -106,15 +105,18 @@ class DataScienceScen(Scenario): competition=self.competition, ) - def get_scenario_all_desc(self) -> str: + def get_scenario_all_desc(self, eda_output=None) -> str: + """ + eda_output depends on dynamic .md files from current workspace, not fixed. + """ return T(".prompts:scenario_description").r( background=self.background, submission_specifications=self.submission_specifications, evaluation=self.metric_description, metric_name=self.metric_name, metric_direction=self.metric_direction, - eda_output=self.eda_output, time_limit=f"{DS_RD_SETTING.full_timeout / 60 / 60 : .2f} hours", + eda_output=eda_output, ) def get_runtime_environment(self) -> str: