Skip to content

Commit a15a06a

Browse files
authored
feat: checkpoint selection (microsoft#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
1 parent 133778c commit a15a06a

File tree

18 files changed

+353
-49
lines changed

18 files changed

+353
-49
lines changed

rdagent/app/data_science/loop.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@
2727
from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
2828
from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
2929
from rdagent.scenarios.data_science.proposal.exp_gen import DSExpGen, DSTrace
30+
from rdagent.scenarios.data_science.proposal.exp_gen.select import (
31+
LatestCKPSelector,
32+
SOTAJumpCKPSelector,
33+
)
3034
from rdagent.scenarios.kaggle.kaggle_crawler import download_data
3135

3236

@@ -49,6 +53,7 @@ def __init__(self, PROP_SETTING: BasePropSetting):
4953

5054
# 2) task generation from a complete solution
5155
# self.exp_gen: ExpGen = import_class(PROP_SETTING.exp_gen)(scen)
56+
self.ckp_selector = LatestCKPSelector()
5257
self.exp_gen = DSExpGen(scen)
5358
self.data_loader_coder = DataLoaderCoSTEER(scen)
5459
self.feature_coder = FeatureCoSTEER(scen)
@@ -68,7 +73,8 @@ def __init__(self, PROP_SETTING: BasePropSetting):
6873
super(RDLoop, self).__init__()
6974

7075
def direct_exp_gen(self, prev_out: dict[str, Any]):
71-
exp = self.exp_gen.gen(self.trace)
76+
selection = self.ckp_selector.get_selection(self.trace)
77+
exp = self.exp_gen.gen(self.trace, selection)
7278
logger.log_object(exp)
7379

7480
# FIXME: this is for LLM debug webapp, remove this when the debugging is done.
@@ -126,6 +132,10 @@ def feedback(self, prev_out: dict[str, Any]) -> ExperimentFeedback:
126132
return feedback
127133

128134
def record(self, prev_out: dict[str, Any]):
135+
136+
# set the DAG parent for the trace
137+
self.trace.sync_dag_parent_and_hist()
138+
129139
e = prev_out.get(self.EXCEPTION_KEY, None)
130140
if e is None:
131141
self.trace.hist.append((prev_out["running"], prev_out["feedback"]))

rdagent/components/coder/data_science/ensemble/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ def implement_one_task(
7474
)
7575

7676
# Generate code with knowledge integration
77-
competition_info = self.scen.get_scenario_all_desc()
77+
competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None))
7878
system_prompt = T(".prompts:ensemble_coder.system").r(
7979
task_desc=ensemble_information_str,
8080
competition_info=competition_info,

rdagent/components/coder/data_science/feature/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def implement_one_task(
6161

6262
# 2. code
6363
system_prompt = T(".prompts:feature_coder.system").r(
64-
competition_info=self.scen.get_scenario_all_desc(),
64+
competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)),
6565
task_desc=feature_information_str,
6666
data_loader_code=workspace.file_dict.get("load_data.py"),
6767
queried_similar_successful_knowledge=queried_similar_successful_knowledge,

rdagent/components/coder/data_science/model/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def implement_one_task(
6262
# 2. code
6363
system_prompt = T(".prompts:model_coder.system").r(
6464
task_desc=model_information_str,
65-
competition_info=self.scen.get_scenario_all_desc(),
65+
competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)),
6666
data_loader_code=workspace.file_dict.get("load_data.py"),
6767
feature_code=workspace.file_dict["feature.py"],
6868
queried_similar_successful_knowledge=queried_similar_successful_knowledge,

rdagent/components/coder/data_science/pipeline/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def implement_one_task(
6666
workspace: FBWorkspace | None = None,
6767
prev_task_feedback: CoSTEERSingleFeedback | None = None,
6868
) -> dict[str, str]:
69-
competition_info = self.scen.get_scenario_all_desc()
69+
competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None))
7070
runtime_environment = self.scen.get_runtime_environment()
7171
data_folder_info = self.scen.processed_data_folder_description
7272
pipeline_task_info = target_task.get_task_information()

rdagent/components/coder/data_science/pipeline/eval.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,8 +119,10 @@ def evaluate(
119119
)
120120
stdout += "\n" + submission_check_out
121121

122+
eda_output = implementation.file_dict.get("EDA.md", None)
123+
122124
system_prompt = T(".prompts:pipeline_eval.system").r(
123-
scenario=self.scen.get_scenario_all_desc(),
125+
scenario=self.scen.get_scenario_all_desc(eda_output=eda_output),
124126
task_desc=target_task.get_task_information(),
125127
spec=T("scenarios.data_science.share:component_spec.Pipeline").r(),
126128
)

rdagent/components/coder/data_science/raw_data_loader/__init__.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ def implement_one_task(
6767
) -> dict[str, str]:
6868
# return a workspace with "load_data.py", "spec/load_data.md" inside
6969
# assign the implemented code to the new workspace.
70-
competition_info = self.scen.get_scenario_all_desc()
70+
competition_info = self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None))
7171
runtime_environment = self.scen.get_runtime_environment()
7272
data_folder_info = self.scen.processed_data_folder_description
7373
data_loader_task_info = target_task.get_task_information()
@@ -231,5 +231,9 @@ def develop(self, exp):
231231
stdout = new_exp.experiment_workspace.execute(env=env, entry=f"python test/data_loader_test.py")
232232
match = re.search(r"(.*?)=== Start of EDA part ===(.*)=== End of EDA part ===", stdout, re.DOTALL)
233233
eda_output = match.groups()[1] if match else None
234-
self.scen.eda_output = eda_output
234+
if eda_output is not None:
235+
new_exp.experiment_workspace.inject_files(**{"EDA.md": eda_output})
236+
else:
237+
eda_output = "No EDA output."
238+
new_exp.experiment_workspace.inject_files(**{"EDA.md": eda_output})
235239
return new_exp

rdagent/components/coder/data_science/workflow/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def implement_one_task(
5959
# 2. code
6060
system_prompt = T(".prompts:workflow_coder.system").r(
6161
task_desc=workflow_information_str,
62-
competition_info=self.scen.get_scenario_all_desc(),
62+
competition_info=self.scen.get_scenario_all_desc(eda_output=workspace.file_dict.get("EDA.md", None)),
6363
queried_similar_successful_knowledge=queried_similar_successful_knowledge,
6464
queried_former_failed_knowledge=queried_former_failed_knowledge[0],
6565
out_spec=PythonAgentOut.get_spec(),

rdagent/components/coder/data_science/workflow/eval.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,8 @@ def evaluate(
127127
stdout += "\n" + submission_check_out
128128

129129
system_prompt = T(".prompts:workflow_eval.system").r(
130-
scenario=self.scen.get_scenario_all_desc(),
130+
# here we pass `None` to `eda_output` because we do not have nor need EDA output for workflow.
131+
scenario=self.scen.get_scenario_all_desc(eda_output=None),
131132
task_desc=target_task.get_task_information(),
132133
spec=(
133134
implementation.file_dict["spec/workflow.md"]

rdagent/core/proposal.py

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
""" """
1+
# TODO: remove `self.scen` if traces will be passed into the instance.
22

33
from __future__ import annotations
44

@@ -112,9 +112,16 @@ def __str__(self) -> str:
112112

113113

114114
class Trace(Generic[ASpecificScen, ASpecificKB]):
115+
NodeType = tuple[Experiment, ExperimentFeedback] # Define NodeType as a new type representing the tuple
116+
115117
def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None:
116118
self.scen: ASpecificScen = scen
117-
self.hist: list[tuple[Experiment, ExperimentFeedback]] = []
119+
self.hist: list[Trace.NodeType] = (
120+
[]
121+
) # List of tuples containing experiments and their feedback, organized over time.
122+
self.dag_parent: list[tuple[int, ...]] = [] # List of tuples representing parent indices in the DAG structure.
123+
# (,) represents no parent; (1,) presents one parent; (1, 2) represents two parents.
124+
118125
# TODO: self.hist is 2-tuple now, remove hypothesis from it, change old code for this later.
119126
self.knowledge_base: ASpecificKB | None = knowledge_base
120127

@@ -128,13 +135,32 @@ def get_sota_hypothesis_and_experiment(self) -> tuple[Hypothesis | None, Experim
128135
return None, None
129136

130137

138+
class CheckpointSelector:
139+
"""
140+
In the trace, we may start from any check point (we'll represent it as a variable `from_checkpoint_idx`)
141+
"""
142+
143+
@abstractmethod
144+
def get_selection(self, trace: Trace) -> tuple[int, ...] | None:
145+
"""
146+
checkpoint_idx represents the place where we want to create a new node.
147+
the return value should be the idx of target node (the parent of the new generating node).
148+
- `(-1, )` represents starting from the latest trial in the trace - default value
149+
- `(idx, )` represents starting from the `idx`-th trial in the trace.
150+
- `None` represents starting from scratch (start a new trace)
151+
152+
153+
- More advanced selection strategies in `select.py`
154+
"""
155+
156+
131157
class ExpGen(ABC):
132158

133159
def __init__(self, scen: Scenario) -> None:
134160
self.scen = scen
135161

136162
@abstractmethod
137-
def gen(self, trace: Trace) -> Experiment:
163+
def gen(self, trace: Trace, selection: tuple[int, ...] = (-1,)) -> Experiment:
138164
"""
139165
Generate the experiment based on the trace.
140166

0 commit comments

Comments
 (0)