feat: add loop_n parameter to the main loop (#611)

RolandMinrui · Xu · web-flow · commit 778c16696225 · 2025-02-18T19:22:52.000+08:00
* add loop_n parameter to main loop

* complete the loop_n

---------

Co-authored-by: Xu &lt;v-xuminrui@microsoft.com&gt;
diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py
@@ -140,7 +140,7 @@ def record(self, prev_out: dict[str, Any]):
         logger.log_object(self.trace.sota_experiment(), tag="SOTA experiment")
 
 
-def main(path=None, step_n=None, competition="bms-molecular-translation"):
+def main(path=None, step_n=None, loop_n=None, competition="bms-molecular-translation"):
     """
 
     Parameters
@@ -149,6 +149,10 @@ def main(path=None, step_n=None, competition="bms-molecular-translation"):
         path like `$LOG_PATH/__session__/1/0_propose`. It indicates that we restore the state that after finish the step 0 in loop1
     step_n :
         How many steps to run; if None, it will run forever until error or KeyboardInterrupt
+    loop_n :
+        How many loops to run; if None, it will run forever until error or KeyboardInterrupt
+        - 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 :
 
 
@@ -174,7 +178,7 @@ def main(path=None, step_n=None, competition="bms-molecular-translation"):
         kaggle_loop = DataScienceRDLoop(DS_RD_SETTING)
     else:
         kaggle_loop = DataScienceRDLoop.load(path)
-    kaggle_loop.run(step_n=step_n)
+    kaggle_loop.run(step_n=step_n, loop_n=loop_n)
 
 
 if __name__ == "__main__":
diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py
@@ -91,21 +91,27 @@ def __init__(self) -> None:
         self.loop_trace = defaultdict(list[LoopTrace])  # the key is the number of loop
         self.session_folder = logger.log_trace_path / "__session__"
 
-    def run(self, step_n: int | None = None) -> None:
+    def run(self, step_n: int | None = None, loop_n: int | None = None) -> None:
         """
 
         Parameters
         ----------
         step_n : int | None
             How many steps to run;
             `None` indicates to run forever until error or KeyboardInterrupt
+        loop_n: int | None
+            How many steps to run; if current loop is incomplete, it will be counted as the first loop for completion
+            `None` indicates to run forever until error or KeyboardInterrupt
         """
         with tqdm(total=len(self.steps), desc="Workflow Progress", unit="step") as pbar:
             while True:
                 if step_n is not None:
                     if step_n <= 0:
                         break
                     step_n -= 1
+                if loop_n is not None:
+                    if loop_n <= 0:
+                        break
 
                 li, si = self.loop_idx, self.step_idx
                 name = self.steps[si]
@@ -141,6 +147,8 @@ def run(self, step_n: int | None = None) -> None:
                 self.step_idx = (self.step_idx + 1) % len(self.steps)
                 if self.step_idx == 0:  # reset to step 0 in next round
                     self.loop_idx += 1
+                    if loop_n is not None:
+                        loop_n -= 1
                     self.loop_prev_out = {}
                     pbar.reset()  # reset the progress bar for the next loop
 
