diff --git a/rdagent/log/ui/ds_trace.py b/rdagent/log/ui/ds_trace.py index dd118756..f59274eb 100644 --- a/rdagent/log/ui/ds_trace.py +++ b/rdagent/log/ui/ds_trace.py @@ -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() diff --git a/rdagent/log/ui/utils.py b/rdagent/log/ui/utils.py index 02372e75..02ac7c60 100644 --- a/rdagent/log/ui/utils.py +++ b/rdagent/log/ui/utils.py @@ -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="%{hovertext}
" + + "Start: %{base}
" + + "End: %{x}
" + + "Duration: %{customdata[0]}
" + + "", + ) + + # 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."),