save more time info (#1151)

Author: XianBW

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"],

rdagent/oai/backend/base.py

@@ -257,16 +257,16 @@ def build_chat_completion(self, user_prompt: str, *args, **kwargs) -> str:  # ty
         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,19 +409,19 @@ def build_messages_and_create_chat_completion(  # type: ignore[no-untyped-def]
             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,
             chat_completion=True,
             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

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 @@ async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
                         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 @@ async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
                     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

rdagent/utils/workflow/loop.py

@@ -10,10 +10,10 @@
 
 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 @@ def __new__(mcs, clsname: str, bases: tuple[type, ...], attrs: dict[str, Any]) -
 
 @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 @@ async def _run_step(self, li: int, force_subproc: bool = False) -> None:
             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 @@ async def _run_step(self, li: int, force_subproc: bool = False) -> None:
                     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: