diff --git a/rdagent/app/qlib_rd_loop/model_w_sc.py b/rdagent/app/qlib_rd_loop/model_w_sc.py index aa6c7b31..c5fca4c3 100644 --- a/rdagent/app/qlib_rd_loop/model_w_sc.py +++ b/rdagent/app/qlib_rd_loop/model_w_sc.py @@ -67,20 +67,20 @@ class ModelLoop(LoopBase, metaclass=LoopMeta): self.trace.hist.append((prev_out["propose"],prev_out["running"] , feedback)) -def main(path=None): +def main(path=None, step_n=None): """ You can continue running session by .. code-block:: python - dotenv run -- python rdagent/app/qlib_rd_loop/model_w_sc.py $LOG_PATH/__session__/1/0_propose + dotenv run -- python rdagent/app/qlib_rd_loop/model_w_sc.py $LOG_PATH/__session__/1/0_propose --step_n 1 # `step_n` is a optional paramter """ if path is None: model_loop = ModelLoop() else: model_loop = ModelLoop.load(path) - model_loop.run() + model_loop.run(step_n=step_n) if __name__ == "__main__": diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py index ee37d989..5a8b1dda 100644 --- a/rdagent/utils/workflow.py +++ b/rdagent/utils/workflow.py @@ -51,9 +51,22 @@ 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): + def run(self, step_n: int | None = None): + """ + + Parameters + ---------- + step_n : int | None + How many steps to run; + `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 + li, si = self.loop_idx, self.step_idx start = datetime.datetime.now(datetime.timezone.utc)