From 2ed6ec53d937dc741ba1303b8bd12ca94617bbca Mon Sep 17 00:00:00 2001
From: Roland Minrui <114476598+RolandMinrui@users.noreply.github.com>
Date: Tue, 4 Mar 2025 15:18:36 +0800
Subject: [PATCH] feat: add do_truncate control for the load function (#656)

feat: add do_truncate control for the load function (#656)
---
 rdagent/app/data_science/loop.py | 10 +++++++---
 rdagent/utils/workflow.py        | 10 +++++++---
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py
index 6cef32f5..2c974025 100644
--- a/rdagent/app/data_science/loop.py
+++ b/rdagent/app/data_science/loop.py
@@ -142,13 +142,15 @@ class DataScienceRDLoop(RDLoop):
         logger.log_object(self.trace.sota_experiment(), tag="SOTA experiment")
 
 
-def main(path=None, output_path=None, step_n=None, loop_n=None, competition="bms-molecular-translation"):
+def main(
+    path=None, output_path=None, step_n=None, loop_n=None, competition="bms-molecular-translation", do_truncate=True
+):
     """
 
     Parameters
     ----------
     path :
-        path like `$LOG_PATH/__session__/1/0_propose`. It indicates that we restore the state that after finish the step 0 in loop1
+        path like `$LOG_PATH/__session__/1/0_propose`. It indicates that we restore the state that after finish the step 0 in loop 1
     output_path :
         path like `$LOG_PATH`. It indicates that where we want to save our session and log information.
     step_n :
@@ -158,6 +160,8 @@ def main(path=None, output_path=None, step_n=None, loop_n=None, competition="bms
         - if current loop is incomplete, it will be counted as the first loop for completion.
         - if both step_n and loop_n are provided, the process will stop as soon as either condition is met.
     competition :
+    do_truncate :
+        If set to True, the logger will truncate the future log messages by calling `logger.storage.truncate`.
 
 
     Auto R&D Evolving loop for models in a Kaggle scenario.
@@ -181,7 +185,7 @@ def main(path=None, output_path=None, step_n=None, loop_n=None, competition="bms
     if path is None:
         kaggle_loop = DataScienceRDLoop(DS_RD_SETTING)
     else:
-        kaggle_loop = DataScienceRDLoop.load(path, output_path)
+        kaggle_loop = DataScienceRDLoop.load(path, output_path, do_truncate)
     kaggle_loop.run(step_n=step_n, loop_n=loop_n)
 
 
diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py
index d187b966..0d35a0ee 100644
--- a/rdagent/utils/workflow.py
+++ b/rdagent/utils/workflow.py
@@ -161,7 +161,9 @@ class LoopBase:
             pickle.dump(self, f)
 
     @classmethod
-    def load(cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None) -> "LoopBase":
+    def load(
+        cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None, do_truncate: bool = True
+    ) -> "LoopBase":
         path = Path(path)
         with path.open("rb") as f:
             session = cast(LoopBase, pickle.load(f))
@@ -175,8 +177,10 @@ class LoopBase:
         # set trace path
         logger.set_trace_path(session.session_folder.parent)
 
-        max_loop = max(session.loop_trace.keys())
-        logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
+        # truncate future message
+        if do_truncate:
+            max_loop = max(session.loop_trace.keys())
+            logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
         return session