From c2bde6dcc3e15953d478f779919cddb3e2f90f92 Mon Sep 17 00:00:00 2001 From: XianBW <36835909+XianBW@users.noreply.github.com> Date: Mon, 4 Aug 2025 17:28:25 +0800 Subject: [PATCH] save more time info (#1151) --- rdagent/log/ui/ds_trace.py | 24 +++++++++---------- rdagent/oai/backend/base.py | 12 +++++----- .../proposal/exp_gen/router/__init__.py | 18 ++++++++++++-- rdagent/utils/workflow/loop.py | 17 +++++++++---- 4 files changed, 45 insertions(+), 26 deletions(-) diff --git a/rdagent/log/ui/ds_trace.py b/rdagent/log/ui/ds_trace.py index 5cbdac00..dd118756 100644 --- a/rdagent/log/ui/ds_trace.py +++ b/rdagent/log/ui/ds_trace.py @@ -534,19 +534,17 @@ def main_win(loop_id, llm_data=None): ) if "running" in loop_data: # get last SOTA_exp_to_submit - current_trace = loop_data["record"]["trace"] - current_selection = current_trace.get_current_selection() - if len(current_selection) > 0: # TODO: Why current_selection can be "()"? - current_idx = current_selection[0] - parent_idxs = current_trace.get_parents(current_idx) - if len(parent_idxs) >= 2 and hasattr(current_trace, "idx2loop_id"): - parent_idx = parent_idxs[-2] - parent_loop_id = current_trace.idx2loop_id[parent_idx] - sota_exp = state.data[parent_loop_id]["record"].get("sota_exp_to_submit", None) - else: - sota_exp = None - else: - sota_exp = None + sota_exp = None + if "record" in loop_data: + current_trace = loop_data["record"]["trace"] + current_selection = current_trace.get_current_selection() + if len(current_selection) > 0: # TODO: Why current_selection can be "()"? + current_idx = current_selection[0] + parent_idxs = current_trace.get_parents(current_idx) + if len(parent_idxs) >= 2 and hasattr(current_trace, "idx2loop_id"): + parent_idx = parent_idxs[-2] + parent_loop_id = current_trace.idx2loop_id[parent_idx] + sota_exp = state.data[parent_loop_id]["record"].get("sota_exp_to_submit", None) running_win( loop_data["running"], diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py index 997db577..4f890231 100644 --- a/rdagent/oai/backend/base.py +++ b/rdagent/oai/backend/base.py @@ -257,16 +257,16 @@ class ChatSession: messages = self.build_chat_completion_message(user_prompt) with logger.tag(f"session_{self.conversation_id}"): - start_time = time.time() + start_time = datetime.now(pytz.timezone("Asia/Shanghai")) response: str = self.api_backend._try_create_chat_completion_or_embedding( # noqa: SLF001 *args, messages=messages, chat_completion=True, **kwargs, ) - end_time = time.time() + end_time = datetime.now(pytz.timezone("Asia/Shanghai")) logger.log_object( - {"user": user_prompt, "resp": response, "duration": end_time - start_time}, tag="debug_llm" + {"user": user_prompt, "resp": response, "start": start_time, "end": end_time}, tag="debug_llm" ) messages.append( @@ -409,7 +409,7 @@ class APIBackend(ABC): shrink_multiple_break=shrink_multiple_break, ) - start_time = time.time() + start_time = datetime.now(pytz.timezone("Asia/Shanghai")) resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] *args, messages=messages, @@ -417,11 +417,11 @@ class APIBackend(ABC): chat_cache_prefix=chat_cache_prefix, **kwargs, ) - end_time = time.time() + end_time = datetime.now(pytz.timezone("Asia/Shanghai")) if isinstance(resp, list): raise ValueError("The response of _try_create_chat_completion_or_embedding should be a string.") logger.log_object( - {"system": system_prompt, "user": user_prompt, "resp": resp, "duration": end_time - start_time}, + {"system": system_prompt, "user": user_prompt, "resp": resp, "start": start_time, "end": end_time}, tag="debug_llm", ) return resp diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py b/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py index 5171ba40..e22b48c1 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/router/__init__.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from datetime import timedelta +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING from rdagent.app.data_science.conf import DS_RD_SETTING @@ -91,12 +91,16 @@ class ParallelMultiTraceExpGen(ExpGen): trace.set_current_selection(local_selection) ds_plan = self.planner.plan(trace) if DS_RD_SETTING.enable_planner else DSExperimentPlan() + + start = datetime.now(timezone.utc) + exp_gen_type = "" if ( (not timer.started or timer.remain_time() >= timedelta(hours=DS_RD_SETTING.merge_hours)) and trace.sota_experiment(selection=local_selection) is None and DS_RD_SETTING.enable_draft_before_first_sota ): exp = self.draft_exp_gen.gen(trace, plan=ds_plan) + exp_gen_type = type(self.draft_exp_gen).__name__ elif ( timer.started and timer.remain_time() < timedelta(hours=DS_RD_SETTING.merge_hours) @@ -105,10 +109,20 @@ class ParallelMultiTraceExpGen(ExpGen): DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000 DS_RD_SETTING.consecutive_errors = 100000 exp = self.merge_exp_gen.gen(trace, plan=ds_plan) + exp_gen_type = type(self.merge_exp_gen).__name__ else: # If there is a sota experiment in the sub-trace and not in merge time, we use default exp_gen exp = self.exp_gen.gen(trace, plan=ds_plan) - + exp_gen_type = type(self.exp_gen).__name__ + end = datetime.now(timezone.utc) + logger.log_object( + { + "exp_gen_type": exp_gen_type, + "start_time": start, + "end_time": end, + }, + tag="exp_gen_time_info", + ) exp.set_local_selection(local_selection) exp.plan = ds_plan return exp diff --git a/rdagent/utils/workflow/loop.py b/rdagent/utils/workflow/loop.py index 94f2d685..5f43a02d 100644 --- a/rdagent/utils/workflow/loop.py +++ b/rdagent/utils/workflow/loop.py @@ -10,10 +10,10 @@ Postscripts: import asyncio import concurrent.futures -import datetime import pickle from collections import defaultdict from dataclasses import dataclass +from datetime import datetime, timezone from pathlib import Path from typing import Any, Callable, Optional, Union, cast @@ -72,8 +72,8 @@ class LoopMeta(type): @dataclass class LoopTrace: - start: datetime.datetime # the start time of the trace - end: datetime.datetime # the end time of the trace + start: datetime # the start time of the trace + end: datetime # the end time of the trace step_idx: int # TODO: more information about the trace @@ -211,7 +211,7 @@ class LoopBase: self.tracker.log_workflow_state() with logger.tag(f"Loop_{li}.{name}"): - start = datetime.datetime.now(datetime.timezone.utc) + start = datetime.now(timezone.utc) func: Callable[..., Any] = cast(Callable[..., Any], getattr(self, name)) next_step_idx = si + 1 @@ -236,8 +236,15 @@ class LoopBase: self.loop_prev_out[li][name] = result # Record the trace - end = datetime.datetime.now(datetime.timezone.utc) + end = datetime.now(timezone.utc) self.loop_trace[li].append(LoopTrace(start, end, step_idx=si)) + logger.log_object( + { + "start_time": start, + "end_time": end, + }, + tag="time_info", + ) # Save snapshot after completing the step self.dump(self.session_folder / f"{li}" / f"{si}_{name}") except Exception as e: