chore: show timeline and change load_times (#1161)

* show timeline, deperacate old load_times * fix bug
2026-07-29 08:27:43 +00:00 · 2025-08-06 18:42:17 +08:00
parent 3244ba6833
commit 0334805744
2 changed files with 256 additions and 64 deletions
@@ -19,7 +19,8 @@ from rdagent.log.ui.conf import UI_SETTING
 from rdagent.log.ui.utils import (
    curve_figure,
    get_sota_exp_stat,
-    load_times,
+    load_times_info,
+    timeline_figure,
    trace_figure,
 )
 from rdagent.log.utils import (
@@ -170,9 +171,11 @@ def load_stdout(stdout_path: Path):

 # UI windows
 def task_win(task):
-    with st.container(border=True):
-        st.markdown(f"**:violet[{task.name}]**")
+    with st.expander(f"**:violet[{task.name}]**", expanded=False):
        st.markdown(task.description)
+        if hasattr(task, "package_info"):
+            st.markdown(f"**:blue[Package Info:]**")
+            st.code(task.package_info)
        if hasattr(task, "architecture"):  # model task
            st.markdown(
                f"""
@@ -185,14 +188,17 @@ def task_win(task):

 def workspace_win(workspace, cmp_workspace=None, cmp_name="last code."):
    show_files = {k: v for k, v in workspace.file_dict.items() if "test" not in k}
-
    if len(show_files) > 0:
        if cmp_workspace:
            diff = generate_diff_from_dict(cmp_workspace.file_dict, show_files, "main.py")
            with st.popover(f":violet[**Diff with {cmp_name}**]", use_container_width=True, icon="🔍"):
                st.code("".join(diff), language="diff", wrap_lines=True, line_numbers=True)
+
+        rtime = workspace.running_info.running_time
+        time_str = timedelta_to_str(timedelta(seconds=rtime) if rtime else None) or "00:00:00"
+
        with st.popover(
-            f"Files in :blue[{replace_ep_path(workspace.workspace_path)}]", use_container_width=True, icon="📂"
+            f"⏱️{time_str} 📂Files in :blue[{replace_ep_path(workspace.workspace_path)}]", use_container_width=True
        ):
            code_tabs = st.tabs(show_files.keys())
            for ct, codename in zip(code_tabs, show_files.keys()):
@@ -276,7 +282,18 @@ def llm_log_win(llm_d: list):
            system = d["obj"].get("system", None)
            user = d["obj"]["user"]
            resp = d["obj"]["resp"]
-            with st.expander(f"**LLM**", icon="🤖", expanded=False):
+            start_time = d["obj"].get("start", "")
+            end_time = d["obj"].get("end", "")
+            if start_time and end_time:
+                start_str = start_time.strftime("%m-%d %H:%M:%S")
+                end_str = end_time.strftime("%m-%d %H:%M:%S")
+                duration = end_time - start_time
+                time_info_str = (
+                    f"🕰️:blue[**{start_str} ~ {end_str}**] ⏳:violet[**{round(duration.total_seconds(), 2)}s**]"
+                )
+            else:
+                time_info_str = ""
+            with st.expander(f"**LLM** {time_info_str}", icon="🤖", expanded=False):
                t1, t2, t3, t4 = st.tabs(
                    [":green[**Response**]", ":blue[**User**]", ":orange[**System**]", ":violet[**ChatBot**]"]
                )
@@ -367,13 +384,13 @@ def exp_gen_win(exp_gen_data, llm_data=None):
    st.header("Exp Gen", divider="blue", anchor="exp-gen")
    if state.show_llm_log and llm_data is not None:
        llm_log_win(llm_data["no_tag"])
-    st.subheader("Hypothesis")
+    st.subheader("💡 Hypothesis")
    hypothesis_win(exp_gen_data["no_tag"].hypothesis)

-    st.subheader("pending_tasks")
+    st.subheader("📋 pending_tasks")
    for tasks in exp_gen_data["no_tag"].pending_tasks_list:
        task_win(tasks[0])
-    st.subheader("Exp Workspace")
+    st.subheader("📁 Exp Workspace")
    workspace_win(exp_gen_data["no_tag"].experiment_workspace)


@@ -573,8 +590,8 @@ def replace_ep_path(p: Path):
 def get_llm_call_stats(llm_data: dict) -> tuple[int, int]:
    total_llm_call = 0
    total_filter_call = 0
-    total_call_seconds = 0
-    filter_call_seconds = 0
+    total_call_duration = timedelta()
+    filter_call_duration = timedelta()
    filter_sys_prompt = T("rdagent.utils.prompts:filter_redundant_text.system").r()
    for li, loop_d in llm_data.items():
        for fn, loop_fn_d in loop_d.items():
@@ -582,12 +599,14 @@ def get_llm_call_stats(llm_data: dict) -> tuple[int, int]:
                for d in v:
                    if "debug_llm" in d["tag"]:
                        total_llm_call += 1
-                        total_call_seconds += d["obj"].get("duration", 0)
+                        total_call_duration += d["obj"].get("end", timedelta()) - d["obj"].get("start", timedelta())
                        if "system" in d["obj"] and filter_sys_prompt == d["obj"]["system"]:
                            total_filter_call += 1
-                            filter_call_seconds += d["obj"].get("duration", 0)
+                            filter_call_duration += d["obj"].get("end", timedelta()) - d["obj"].get(
+                                "start", timedelta()
+                            )

-    return total_llm_call, total_filter_call, total_call_seconds, filter_call_seconds
+    return total_llm_call, total_filter_call, total_call_duration, filter_call_duration


 def get_timeout_stats(llm_data: dict):
@@ -638,13 +657,13 @@ def summarize_win():
        with info3.popover("RDLOOP", icon="⚙️"):
            st.write(state.data.get("settings", {}).get("RDLOOP_SETTINGS", "No settings found."))

-        llm_call, llm_filter_call, llm_call_seconds, llm_filter_call_seconds = get_llm_call_stats(state.llm_data)
-        info4.metric("LLM Calls", llm_call, help=timedelta_to_str(timedelta(seconds=llm_call_seconds)))
+        llm_call, llm_filter_call, llm_call_duration, filter_call_duration = get_llm_call_stats(state.llm_data)
+        info4.metric("LLM Calls", llm_call, help=timedelta_to_str(llm_call_duration))
        info5.metric(
            "LLM Filter Calls",
            llm_filter_call,
            delta=-round(llm_filter_call / llm_call, 5),
-            help=timedelta_to_str(timedelta(seconds=llm_filter_call_seconds)),
+            help=timedelta_to_str(filter_call_duration),
        )

        timeout_stats = get_timeout_stats(state.llm_data)
@@ -718,21 +737,28 @@ def summarize_win():
            df.loc[loop, "COST($)"] = sum(tc.content["cost"] for tc in state.token_costs[loop])

            # Time Stats
-            if loop in state.times and state.times[loop]:
-                exp_gen_time = coding_time = running_time = None
-                all_steps_time = timedelta()
-                for lpt in state.times[loop]:
-                    all_steps_time += lpt.end - lpt.start
-                    if lpt.step_idx == 0:
-                        exp_gen_time = lpt.end - lpt.start
-                    elif lpt.step_idx == 1:
-                        coding_time = lpt.end - lpt.start
-                    elif lpt.step_idx == 2:
-                        running_time = lpt.end - lpt.start
-                df.loc[loop, "Time"] = timedelta_to_str(all_steps_time)
-                df.loc[loop, "Exp Gen"] = timedelta_to_str(exp_gen_time)
-                df.loc[loop, "Coding"] = timedelta_to_str(coding_time)
-                df.loc[loop, "Running"] = timedelta_to_str(running_time)
+            exp_gen_time = timedelta()
+            coding_time = timedelta()
+            running_time = timedelta()
+            all_steps_time = timedelta()
+            if loop in state.times:
+                for step_name, step_time in state.times[loop].items():
+                    step_duration = step_time["end_time"] - step_time["start_time"]
+                    if step_name == "exp_gen":
+                        exp_gen_time += step_duration
+                        all_steps_time += step_duration
+                    elif step_name == "coding":
+                        coding_time += step_duration
+                        all_steps_time += step_duration
+                    elif step_name == "running":
+                        running_time += step_duration
+                        all_steps_time += step_duration
+                    elif step_name in ["feedback", "record"]:
+                        all_steps_time += step_duration
+            df.loc[loop, "Time"] = timedelta_to_str(all_steps_time)
+            df.loc[loop, "Exp Gen"] = timedelta_to_str(exp_gen_time)
+            df.loc[loop, "Coding"] = timedelta_to_str(coding_time)
+            df.loc[loop, "Running"] = timedelta_to_str(running_time)

            if "running" in loop_data and "no_tag" in loop_data["running"]:
                try:
@@ -835,22 +861,10 @@ def summarize_win():
            df = df[df["Feedback"] == "✅"]
        st.dataframe(df[df.columns[~df.columns.isin(["Hypothesis", "Reason", "Others"])]])

-        # COST curve
-        costs = df["COST($)"].astype(float)
-        costs.index = [f"L{i}" for i in costs.index]
-        cumulative_costs = costs.cumsum()
-        with st.popover("COST Curve", icon="💰", use_container_width=True):
-            fig = px.line(
-                x=costs.index,
-                y=[costs.values, cumulative_costs.values],
-                labels={"x": "Loop", "value": "COST($)"},
-                title="COST($) per Loop & Cumulative COST($)",
-                markers=True,
-            )
-            fig.update_traces(mode="lines+markers")
-            fig.data[0].name = "COST($) per Loop"
-            fig.data[1].name = "Cumulative COST($)"
-            st.plotly_chart(fig)
+        # timeline figure
+        if state.times:
+            with st.popover("Timeline", icon="⏱️", use_container_width=True):
+                st.plotly_chart(timeline_figure(state.times))

        # scores curve
        vscores = {}
@@ -920,8 +934,8 @@ def summarize_win():
        comp_df["Valid Rate"] = comp_df["Valid Rate"].apply(lambda x: f"{x}%")
        comp_df["Success Rate"] = comp_df["Success Rate"].apply(lambda x: f"{x}%")
        comp_df.loc["Total", "Avg e-loops(c)"] = round(df["e-loops(c)"].mean(), 2)
-        st2.markdown("### Component Statistics")
-        st2.dataframe(comp_df)
+        with st2.popover("Component Statistics", icon="📊", use_container_width=True):
+            st.dataframe(comp_df)

        # component time statistics
        time_df = df.loc[:, ["Component", "Time", "Exp Gen", "Coding", "Running"]]
@@ -933,7 +947,6 @@ def summarize_win():
                "Running": "timedelta64[ns]",
            }
        )
-        st1.markdown("### Time Statistics")
        time_stat_df = time_df.groupby("Component").sum()
        time_stat_df.loc["Total"] = time_stat_df.sum()
        time_stat_df.loc[:, "Exp Gen(%)"] = (time_stat_df["Exp Gen"] / time_stat_df["Time"] * 100).round(2)
@@ -941,7 +954,25 @@ def summarize_win():
        time_stat_df.loc[:, "Running(%)"] = (time_stat_df["Running"] / time_stat_df["Time"] * 100).round(2)
        for col in ["Time", "Exp Gen", "Coding", "Running"]:
            time_stat_df[col] = time_stat_df[col].map(timedelta_to_str)
-        st1.dataframe(time_stat_df)
+        with st1.popover("Time Statistics", icon="⏱️", use_container_width=True):
+            st.dataframe(time_stat_df)
+
+        # COST curve
+        costs = df["COST($)"].astype(float)
+        costs.index = [f"L{i}" for i in costs.index]
+        cumulative_costs = costs.cumsum()
+        with st.popover("COST Curve", icon="💰", use_container_width=True):
+            fig = px.line(
+                x=costs.index,
+                y=[costs.values, cumulative_costs.values],
+                labels={"x": "Loop", "value": "COST($)"},
+                title="COST($) per Loop & Cumulative COST($)",
+                markers=True,
+            )
+            fig.update_traces(mode="lines+markers")
+            fig.data[0].name = "COST($) per Loop"
+            fig.data[1].name = "Cumulative COST($)"
+            st.plotly_chart(fig)


 def stdout_win(loop_id: int):
@@ -1029,7 +1060,7 @@ with st.sidebar:
                st.toast("Please select a log path first!", icon="🟡")
                st.stop()

-            state.times = load_times(state.log_folder / state.log_path)
+            state.times = load_times_info(state.log_folder / state.log_path)
            state.data, state.llm_data, state.token_costs = load_data(state.log_folder / state.log_path)
            state.sota_info = get_sota_exp_stat(Path(state.log_folder) / state.log_path, to_submit=True)
            st.rerun()
@@ -1,13 +1,15 @@
 import math
 import pickle
 import re
-from collections import deque
+from collections import defaultdict, deque
 from datetime import datetime, timedelta
 from pathlib import Path
+from typing import Literal

 import matplotlib.pyplot as plt
 import networkx as nx
 import pandas as pd
+import plotly.express as px
 import plotly.graph_objects as go
 import typer
 from matplotlib import pyplot as plt
@@ -17,7 +19,7 @@ from rdagent.core.proposal import Trace
 from rdagent.core.utils import cache_with_pickle
 from rdagent.log.storage import FileStorage
 from rdagent.log.ui.conf import UI_SETTING
-from rdagent.log.utils import extract_json
+from rdagent.log.utils import extract_json, extract_loopid_func_name
 from rdagent.oai.llm_utils import md5_hash
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.kaggle.kaggle_crawler import get_metric_direction
@@ -230,7 +232,7 @@ def get_sota_exp_stat(


@cache_with_pickle(_log_path_hash_func, force=True)
-def load_times(log_path: Path):
+def load_times_deprecated(log_path: Path):
    try:
        session_path = log_path / "__session__"
        max_li = max(int(p.name) for p in session_path.iterdir() if p.is_dir() and p.name.isdigit())
@@ -243,6 +245,44 @@ def load_times(log_path: Path):
    return rd_times


+@cache_with_pickle(_log_path_hash_func, force=True)
+def load_times_info(log_path: Path) -> dict[int, dict[str, dict[Literal["start_time", "end_time"], datetime]]]:
+    """
+    Load timing information for each loop and step.
+
+    Returns
+    -------
+    dict[int, dict[str, dict[Literal["start_time", "end_time"], datetime]]]
+        Dictionary with loop IDs as keys, where each value contains step names
+        mapping to their start and end times.
+
+        Example:
+            {
+                1: {
+                    "exp_gen": {
+                        "start_time": datetime(2024, 1, 1, 10, 0, 0),
+                        "end_time": datetime(2024, 1, 1, 10, 15, 30)
+                    },
+                    "coding": {
+                        "start_time": datetime(2024, 1, 1, 10, 15, 30),
+                        "end_time": datetime(2024, 1, 1, 10, 45, 12)
+                    }
+                },
+            }
+    """
+    log_storage = FileStorage(log_path)
+    time_msgs = list(log_storage.iter_msg(tag="time_info"))
+    exp_gen_time_msgs = list(log_storage.iter_msg(tag="exp_gen_time_info"))
+    times_info = defaultdict(dict)
+    for msg in time_msgs:
+        li, fn = extract_loopid_func_name(msg.tag)
+        times_info[int(li)][fn] = msg.content
+    for msg in exp_gen_time_msgs:
+        li, fn = extract_loopid_func_name(msg.tag)
+        times_info[int(li)]["exp_gen"] = msg.content
+    return times_info
+
+
 def _log_folders_summary_hash_func(log_folders: list[str], hours: int | None = None):
    hash_str = ""
    for lf in log_folders:
@@ -300,14 +340,21 @@ def get_summary_df(log_folders: list[str], hours: int | None = None) -> tuple[di
            coding_time = timedelta()
            running_time = timedelta()
            all_time = timedelta()
-            times_info = load_times(Path(lf) / k)
-            for time_info in times_info.values():
-                all_time += sum((ti.end - ti.start for ti in time_info), timedelta())
-                exp_gen_time += time_info[0].end - time_info[0].start
-                if len(time_info) > 1:
-                    coding_time += time_info[1].end - time_info[1].start
-                if len(time_info) > 2:
-                    running_time += time_info[2].end - time_info[2].start
+            times_info = load_times_info(Path(lf) / k)
+            for loop_times in times_info.values():
+                for step_name, step_time in loop_times.items():
+                    step_duration = step_time["end_time"] - step_time["start_time"]
+                    if step_name == "exp_gen":
+                        exp_gen_time += step_duration
+                        all_time += step_duration
+                    elif step_name == "coding":
+                        coding_time += step_duration
+                        all_time += step_duration
+                    elif step_name == "running":
+                        all_time += step_duration
+                        running_time += step_duration
+                    elif step_name in ["feedback", "record"]:
+                        all_time += step_duration
            v["exec_time"] = str(all_time).split(".")[0]
            v["exp_gen_time"] = str(exp_gen_time).split(".")[0]
            v["coding_time"] = str(coding_time).split(".")[0]
@@ -791,6 +838,120 @@ def trace_figure(trace: Trace, merge_loops: list = []):
    return fig


+def timeline_figure(times_dict: dict[int, dict[str, dict[Literal["start_time", "end_time"], datetime]]]) -> go.Figure:
+    # Prepare data for px.timeline
+    timeline_data = []
+    step_names = ["exp_gen", "coding", "running", "feedback", "record"]
+
+    # Beautiful color palette with gradients
+    colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#FFA726", "#5A0069"]
+    color_map = {step: color for step, color in zip(step_names, colors)}
+
+    for loop_id, steps in times_dict.items():
+        for step_name, timing in steps.items():
+            if step_name in step_names:
+                duration = timing["end_time"] - timing["start_time"]
+                timeline_data.append(
+                    {
+                        "Start": timing["start_time"],
+                        "Finish": timing["end_time"],
+                        "Step": step_name,
+                        "Loop_ID": f"Loop {loop_id}",
+                        "Duration": str(duration).split(".")[0],  # Remove microseconds
+                    }
+                )
+
+    # Create DataFrame and sort by loop ID in descending order
+    df = pd.DataFrame(timeline_data)
+    df["loop_sort"] = df["Loop_ID"].str.extract("(\d+)").astype(int)
+    df = df.sort_values("loop_sort", ascending=False)
+
+    # Create timeline with enhanced styling
+    fig = px.timeline(
+        df,
+        x_start="Start",
+        x_end="Finish",
+        y="Loop_ID",
+        color="Step",
+        color_discrete_map=color_map,
+        title="🚀 Data Science Loop Timeline",
+        hover_data={"Duration": True, "Loop_ID": False, "Step": False},
+        hover_name="Step",
+    )
+
+    # Enhanced styling and layout
+    fig.update_traces(
+        marker=dict(line=dict(width=1, color="rgba(255,255,255,0.8)"), opacity=0.85),
+        width=0.9,  # Increased from 0.8 to make bars thicker and reduce spacing
+        hovertemplate="<b>%{hovertext}</b><br>"
+        + "Start: %{base}<br>"
+        + "End: %{x}<br>"
+        + "Duration: %{customdata[0]}<br>"
+        + "<extra></extra>",
+    )
+
+    # Beautiful layout with gradients and shadows
+    fig.update_layout(
+        title=dict(text="Data Science Loop Timeline", x=0.0, font=dict(size=24, color="#2C3E50", family="Arial Black")),
+        xaxis=dict(
+            title="⏰ Time",
+            showgrid=True,
+            gridwidth=1,
+            gridcolor="rgba(176, 196, 222, 0.4)",
+            zeroline=False,
+            tickfont=dict(size=12, color="#34495E"),
+            title_font=dict(size=14, color="#2C3E50", family="Arial"),
+        ),
+        yaxis=dict(
+            title="🔄 Loop ID",
+            showgrid=True,
+            gridwidth=1,
+            gridcolor="rgba(176, 196, 222, 0.4)",
+            zeroline=False,
+            tickfont=dict(size=12, color="#34495E"),
+            title_font=dict(size=14, color="#2C3E50", family="Arial"),
+        ),
+        plot_bgcolor="rgba(248, 249, 250, 0.8)",
+        paper_bgcolor="white",
+        height=max(200, len(times_dict) * 25),  # Reduced from 300 and 30 to 200 and 25
+        margin=dict(l=100, r=60, t=80, b=60),
+        legend=dict(
+            x=0.98,
+            y=0.98,
+            xanchor="right",
+            yanchor="top",
+            bgcolor="rgba(255,255,255,0.9)",
+            bordercolor="rgba(0,0,0,0.2)",
+            borderwidth=1,
+            title_font=dict(size=12, color="#2C3E50"),
+            font=dict(size=11, color="#34495E"),
+            traceorder="normal",
+        ),
+        font=dict(family="Arial, sans-serif"),
+        template="plotly_white",
+    )
+
+    # Reorder legend to match step_names order
+    fig.data = sorted(
+        fig.data, key=lambda trace: step_names.index(trace.name) if trace.name in step_names else len(step_names)
+    )
+
+    # Add subtle shadow effect
+    fig.add_shape(
+        type="rect",
+        xref="paper",
+        yref="paper",
+        x0=0,
+        y0=0,
+        x1=1,
+        y1=1,
+        line=dict(color="rgba(0,0,0,0.1)", width=2),
+        fillcolor="rgba(0,0,0,0.02)",
+    )
+
+    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."),