add trace DAG in trace, make exp show compatible with old versions (#1018)

rdagent/log/ui/ds_trace.py

@@ -13,7 +13,7 @@ from streamlit import session_state as state
 
 from rdagent.app.data_science.loop import DataScienceRDLoop
 from rdagent.log.storage import FileStorage
-from rdagent.log.ui.utils import load_times
+from rdagent.log.ui.utils import load_times, trace_figure
 from rdagent.log.utils import (
     LogColors,
     extract_evoid,
@@ -43,7 +43,7 @@ def convert_defaultdict_to_dict(d):
     return d
 
 
-@st.cache_data(persist=True)
+# @st.cache_data(persist=True)
 def load_data(log_path: Path):
     data = defaultdict(lambda: defaultdict(dict))
     llm_data = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
@@ -374,7 +374,10 @@ def running_win(data, base_exp, llm_data=None, sota_exp=None):
             cmp_name="last SOTA",
         )
         st.subheader("Result")
-        st.write(data["no_tag"].result)
+        try:
+            st.write(data["no_tag"].result)
+        except AttributeError as e:  # Compatible with old versions
+            st.write(data["no_tag"].__dict__["result"])
         mle_score_text = data.get("mle_score", "no submission to score")
         mle_score = extract_json(mle_score_text)
         st.subheader(
@@ -458,6 +461,13 @@ def replace_ep_path(p: Path):
 def summarize_data():
     st.header("Summary", divider="rainbow")
     with st.container(border=True):
+        min_id, max_id = get_state_data_range(state.data)
+        if st.toggle("Show trace DAG", key="show_trace_dag"):
+            st.markdown("### Trace DAG")
+            final_trace_loop_id = max_id
+            while "record" not in state.data[final_trace_loop_id]:
+                final_trace_loop_id -= 1
+            st.pyplot(trace_figure(state.data[final_trace_loop_id]["record"]["trace"]))
         df = pd.DataFrame(
             columns=[
                 "Component",
@@ -475,10 +485,9 @@ def summarize_data():
                 "Start Time (UTC+8)",
                 "End Time (UTC+8)",
             ],
-            index=range(len(state.data) - 1),
+            index=range(min_id, max_id + 1),
         )
 
-        min_id, max_id = get_state_data_range(state.data)
         for loop in range(min_id, max_id + 1):
             loop_data = state.data[loop]
             df.loc[loop, "Component"] = loop_data["direct_exp_gen"]["no_tag"].hypothesis.component
@@ -503,9 +512,11 @@ def summarize_data():
                 df.loc[loop, "End Time (UTC+8)"] = state.times[loop][-1].end + timedelta(hours=8)
             if "running" in loop_data and "no_tag" in loop_data["running"]:
                 try:
-                    df.loc[loop, "Running Score (valid)"] = str(
-                        round(loop_data["running"]["no_tag"].result.loc["ensemble"].iloc[0], 5)
-                    )
+                    try:
+                        running_result = loop_data["running"]["no_tag"].result
+                    except AttributeError as e:  # Compatible with old versions
+                        running_result = loop_data["running"]["no_tag"].__dict__["result"]
+                    df.loc[loop, "Running Score (valid)"] = str(round(running_result.loc["ensemble"].iloc[0], 5))
                 except:
                     df.loc[loop, "Running Score (valid)"] = "❌"
                 if "mle_score" not in state.data[loop]:
@@ -720,6 +731,7 @@ with st.sidebar:
     st.toggle("*Show save workspace feature*", key="show_save_input")
     st.markdown(
         f"""
+- [Summary](#summary)
 - [Exp Gen](#exp-gen)
 - [Coding](#coding)
 - [Running](#running)
@@ -732,7 +744,7 @@ with st.sidebar:
 def get_state_data_range(state_data):
     # we have a "competition" key in state_data
     # like dict_keys(['competition', 10, 11, 12, 13, 14])
-    keys = [k for k in state_data.keys() if isinstance(k, int)]
+    keys = [k for k in state_data.keys() if isinstance(k, int) and "no_tag" in state_data[k]["direct_exp_gen"]]
     return min(keys), max(keys)
 
 

rdagent/log/ui/utils.py

@@ -5,10 +5,13 @@ from collections import deque
 from datetime import datetime, timedelta
 from pathlib import Path
 
+import matplotlib.pyplot as plt
+import networkx as nx
 import pandas as pd
 import typer
 
 from rdagent.app.data_science.loop import DataScienceRDLoop
+from rdagent.core.proposal import Trace
 from rdagent.core.utils import cache_with_pickle
 from rdagent.log.ui.conf import UI_SETTING
 from rdagent.log.utils import extract_json
@@ -274,10 +277,16 @@ def get_summary_df(log_folders: list[str], hours: int | None = None) -> tuple[di
             v["running_time"] = str(running_time).split(".")[0]
 
             final_sota_exp = get_final_sota_exp(Path(lf) / k)
-            if final_sota_exp is not None and final_sota_exp.result is not None:
-                v["sota_exp_score_valid"] = final_sota_exp.result.loc["ensemble"].iloc[0]
-            else:
-                v["sota_exp_score_valid"] = None
+
+            v["sota_exp_score_valid"] = None
+            if final_sota_exp is not None:
+                try:
+                    final_sota_result = final_sota_exp.result
+                except AttributeError:  # Compatible with old versions
+                    final_sota_result = final_sota_exp.__dict__["result"]
+                if final_sota_result is not None:
+                    v["sota_exp_score_valid"] = final_sota_result.loc["ensemble"].iloc[0]
+
             v["sota_exp_stat_new"] = get_sota_exp_stat(Path(lf) / k)
             # change experiment name
             if "amlt" in lf:
@@ -543,6 +552,57 @@ def get_statistics_df(summary_df: pd.DataFrame) -> pd.DataFrame:
     return stat_df
 
 
+def trace_figure(trace: Trace):
+    G = nx.DiGraph()
+
+    # Calculate the number of ancestors for each node (root node is 0, more ancestors means lower level)
+    levels = {}
+    for i in range(len(trace.dag_parent)):
+        levels[i] = len(trace.get_parents(i))
+
+    # 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}"))
+    G.add_edges_from(edges)
+
+    # 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}")
+
+    # Layout by level: y axis is -lvl, x axis is evenly distributed
+    pos = {}
+
+    def parent_avg_pos(node):
+        id = int(node[1:])
+        parents = trace.dag_parent[id]
+
+        if not parents:
+            return 0
+
+        parent_nodes = [f"L{p}" for p in parents]
+        parent_xs = [pos[p][0] for p in parent_nodes if p in pos]
+        return sum(parent_xs) / len(parent_xs)
+
+    for lvl in sorted(layer_nodes):
+        nodes = layer_nodes[lvl]
+        # For root nodes, sort directly by index
+        if lvl == 0:
+            sorted_nodes = sorted(nodes, key=lambda n: int(n[1:]))
+        else:
+            sorted_nodes = sorted(nodes, key=parent_avg_pos)
+        y = -lvl
+        x_start = -0.5 * (len(sorted_nodes) - 1)
+        for i, node in enumerate(sorted_nodes):
+            pos[node] = (x_start + i, y)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    nx.draw(G, pos, with_labels=True, arrows=True, node_color="skyblue", node_size=100, font_size=5, ax=ax)
+    return fig
+
+
 def compare(
     exp_list: list[str] = typer.Option(..., "--exp-list", help="List of experiment names.", show_default=False),
     output: str = typer.Option("merge_base_df.h5", help="Output summary file name."),