feat: add output_path to load function of LoopBase (#628)

* feat: add output_path to load from checkpoint function

* add default value to output_path

* sort import

* sort imports

---------

Co-authored-by: Xu <v-xuminrui@microsoft.com>

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, 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"):
     """
 
     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
+    output_path :
+        path like `$LOG_PATH`. It indicates that where we want to save our session and log information.
     step_n :
         How many steps to run; if None, it will run forever until error or KeyboardInterrupt
     loop_n :
@@ -179,7 +181,7 @@ def main(path=None, step_n=None, loop_n=None, competition="bms-molecular-transla
     if path is None:
         kaggle_loop = DataScienceRDLoop(DS_RD_SETTING)
     else:
-        kaggle_loop = DataScienceRDLoop.load(path)
+        kaggle_loop = DataScienceRDLoop.load(path, output_path)
     kaggle_loop.run(step_n=step_n, loop_n=loop_n)
 
 

rdagent/utils/workflow.py

@@ -14,7 +14,7 @@ import time
 from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, TypeVar, cast
+from typing import Any, Callable, Optional, TypeVar, Union, cast
 
 from tqdm.auto import tqdm
 
@@ -161,10 +161,18 @@ class LoopBase:
             pickle.dump(self, f)
 
     @classmethod
-    def load(cls, path: str | Path) -> "LoopBase":
+    def load(cls, path: Union[str, Path], output_path: Optional[Union[str, Path]] = None) -> "LoopBase":
         path = Path(path)
         with path.open("rb") as f:
             session = cast(LoopBase, pickle.load(f))
+
+        # set session folder
+        if output_path:
+            output_path = Path(output_path)
+            output_path.mkdir(parents=True, exist_ok=True)
+            session.session_folder = output_path / "__session__"
+
+        # set trace path
         logger.set_trace_path(session.session_folder.parent)
 
         max_loop = max(session.loop_trace.keys())