gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-01 17:37:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* add loop_n parameter to main loop

* complete the loop_n

---------

Co-authored-by: Xu <v-xuminrui@microsoft.com>
This commit is contained in:
Roland Minrui
2025-02-18 19:22:52 +08:00
committed by GitHub
parent 15584bad43
commit 2b89fafa56
2 changed files with 15 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/utils/workflow.py
+9 -1
View File
@@ -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