From 810af369e0ebdcb97016966cdbcc3e542dbc9fd3 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Mon, 21 Jul 2025 11:02:37 +0800 Subject: [PATCH] feat: add loop ID mapping to trace nodes and update UI labels (#1098) * Add loop index * feat: add loop ID mapping to trace nodes and update UI labels * lint * doc lint --- rdagent/core/proposal.py | 24 ++++++++++++++++++- rdagent/log/ui/utils.py | 15 +++++++++--- rdagent/scenarios/data_science/loop.py | 8 +++++-- .../data_science/proposal/exp_gen/base.py | 16 ++++--------- rdagent/utils/workflow/loop.py | 8 ++++++- 5 files changed, 53 insertions(+), 18 deletions(-) diff --git a/rdagent/core/proposal.py b/rdagent/core/proposal.py index d73ffea2..c05df680 100644 --- a/rdagent/core/proposal.py +++ b/rdagent/core/proposal.py @@ -127,11 +127,29 @@ class Trace(Generic[ASpecificScen, ASpecificKB]): def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None: self.scen: ASpecificScen = scen + + # BEGIN: graph structure ------------------------- 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. + # Definition: + # - (,) represents no parent (root node in one tree); + # - (1,) presents one parent; + # - (1, 2) represents two parents (Multiple parent is not implemented yet). + # Syntax sugar for the parent relationship: + # - Only for selection: + # - (-1,) indicates that select the last record node as parent. + + # NOTE: the sequence of hist and dag_parent is organized by the order to record the experiment. + # So it may be different from the order of the loop_id. + # So we need an extra mapping to map the enqueue id back to the loop id. + self.idx2loop_id: dict[int, int] = {} + + # Design discussion: + # - If we unifiy the loop_id and the enqueue id, we will have less recognition burden. + # - If we use different id for loop and enqueue, we don't have to handle the placeholder logic. + # END: graph structure ------------------------- # TODO: self.hist is 2-tuple now, remove hypothesis from it, change old code for this later. self.knowledge_base: ASpecificKB | None = knowledge_base @@ -227,9 +245,13 @@ class CheckpointSelector: 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 + + - NOTE: we don't encourage to use this option; It is confusing when we have multiple traces. + - `(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` """ diff --git a/rdagent/log/ui/utils.py b/rdagent/log/ui/utils.py index 99b6a955..ade9833c 100644 --- a/rdagent/log/ui/utils.py +++ b/rdagent/log/ui/utils.py @@ -606,13 +606,22 @@ def trace_figure(trace: Trace): for i in range(len(trace.dag_parent)): levels[i] = len(trace.get_parents(i)) + def get_display_name(idx: int): + """ + Convert to index in the queue (enque id) to loop_idx for easier understanding. + """ + if hasattr(trace, "idx2loop_id"): + # FIXME: only keep me after it is stable. Just for compatibility. + return f"L{trace.idx2loop_id[idx]}" + return f"L{idx}" + # Add nodes and edges edges = [] for i, parents in enumerate(trace.dag_parent): for parent in parents: - edges.append((f"L{parent}", f"L{i}")) + edges.append((get_display_name(parent), get_display_name(i))) if len(parents) == 0: - G.add_node(f"L{i}") + G.add_node(get_display_name(i)) G.add_edges_from(edges) # Check if G is a path (a single line) @@ -643,7 +652,7 @@ def trace_figure(trace: Trace): # Group nodes by number of ancestors, fewer ancestors are higher up layer_nodes = {} for idx, lvl in levels.items(): - layer_nodes.setdefault(lvl, []).append(f"L{idx}") + layer_nodes.setdefault(lvl, []).append(get_display_name(idx)) # Layout by level: y axis is -lvl, x axis is evenly distributed pos = {} diff --git a/rdagent/scenarios/data_science/loop.py b/rdagent/scenarios/data_science/loop.py index 690a73ba..5ff1fee1 100644 --- a/rdagent/scenarios/data_science/loop.py +++ b/rdagent/scenarios/data_science/loop.py @@ -200,6 +200,8 @@ class DataScienceRDLoop(RDLoop): exp: DSExperiment = None + cur_loop_id = prev_out[self.LOOP_IDX_KEY] + e = prev_out.get(self.EXCEPTION_KEY, None) if e is None: exp = prev_out["running"] @@ -210,18 +212,20 @@ class DataScienceRDLoop(RDLoop): # set the local selection to the trace as global selection, then set the DAG parent for the trace if exp.local_selection is not None: self.trace.set_current_selection(exp.local_selection) - self.trace.sync_dag_parent_and_hist((exp, prev_out["feedback"])) + self.trace.sync_dag_parent_and_hist((exp, prev_out["feedback"]), cur_loop_id) else: exp: DSExperiment = prev_out["direct_exp_gen"] if isinstance(e, CoderError) else prev_out["coding"] # set the local selection to the trace as global selection, then set the DAG parent for the trace if exp.local_selection is not None: self.trace.set_current_selection(exp.local_selection) + self.trace.sync_dag_parent_and_hist( ( exp, ExperimentFeedback.from_exception(e), - ) + ), + cur_loop_id, ) if self.trace.sota_experiment() is None: diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/base.py b/rdagent/scenarios/data_science/proposal/exp_gen/base.py index 6f8076a1..1e9001f3 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/base.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/base.py @@ -49,18 +49,10 @@ 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. + super().__init__(scen, knowledge_base) - self.knowledge_base = knowledge_base - self.current_selection: tuple[int, ...] = (-1,) + # NOTE: this line is just for linting. + self.hist: list[tuple[DSExperiment, ExperimentFeedback] | None] = [] self.sota_exp_to_submit: DSExperiment | None = None # grab the global best exp to submit @@ -95,6 +87,7 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): def sync_dag_parent_and_hist( self, exp_and_fb: tuple[Experiment, ExperimentFeedback], + cur_loop_id: int, ) -> None: """ Adding corresponding parent index to the dag_parent when the hist is going to be changed. @@ -114,6 +107,7 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]): self.dag_parent.append((current_node_idx,)) self.hist.append(exp_and_fb) + self.idx2loop_id[len(self.hist) - 1] = cur_loop_id def retrieve_search_list( self, diff --git a/rdagent/utils/workflow/loop.py b/rdagent/utils/workflow/loop.py index 683bfa07..24a31ad8 100644 --- a/rdagent/utils/workflow/loop.py +++ b/rdagent/utils/workflow/loop.py @@ -97,6 +97,7 @@ class LoopBase: ] = () # you can define a list of error that will withdraw current loop EXCEPTION_KEY = "_EXCEPTION" + LOOP_IDX_KEY = "_LOOP_IDX" SENTINEL = -1 _pbar: tqdm # progress bar instance @@ -113,7 +114,9 @@ class LoopBase: self.step_idx: defaultdict[int, int] = defaultdict(int) # dict from loop index to next step index self.queue: asyncio.Queue[Any] = asyncio.Queue() - # Store step results for all loops in a nested dictionary: loop_prev_out[loop_index][step_name] + # Store step results for all loops in a nested dictionary, following information will be stored: + # - loop_prev_out[loop_index][step_name]: the output of the step function + # - loop_prev_out[loop_index][]: the special keys self.loop_prev_out: dict[int, dict[str, Any]] = defaultdict(dict) self.loop_trace = defaultdict(list[LoopTrace]) # the key is the number of loop self.session_folder = Path(LOG_SETTINGS.trace_path) / "__session__" @@ -213,6 +216,9 @@ class LoopBase: next_step_idx = si + 1 step_forward = True + # NOTE: each step are aware are of current loop index + # It is very important to set it before calling the step function! + self.loop_prev_out[li][self.LOOP_IDX_KEY] = li try: # Call function with current loop's output, await if coroutine or use ProcessPoolExecutor for sync if required if force_subproc: