diff --git a/rdagent/app/data_science/loop.py b/rdagent/app/data_science/loop.py index 5a5e5084..f248dbc5 100644 --- a/rdagent/app/data_science/loop.py +++ b/rdagent/app/data_science/loop.py @@ -140,7 +140,7 @@ class DataScienceRDLoop(RDLoop): 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 index 9e82d7a4..eb0f8be6 100644 --- a/rdagent/utils/workflow.py +++ b/rdagent/utils/workflow.py @@ -91,7 +91,7 @@ class LoopBase: 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 @@ -99,6 +99,9 @@ class LoopBase: 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: @@ -106,6 +109,9 @@ class LoopBase: 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 @@ class LoopBase: 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